From 20e7b9e9aed2a6a66b864d3b51578950c45c4caa Mon Sep 17 00:00:00 2001
From: alexcbb <alexchapin@hotmail.fr>
Date: Mon, 24 Jul 2023 10:58:55 +0200
Subject: [PATCH] Fix config issue
---
train_sa.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/train_sa.py b/train_sa.py
index 14c0972..5a3f9a7 100644
--- a/train_sa.py
+++ b/train_sa.py
@@ -97,7 +97,7 @@ def main():
shuffle=True, worker_init_fn=data.worker_init_fn)
#### Create model
- model = SlotAttentionAutoEncoder(resolution, num_slots, num_iterations).to(device)
+ model = SlotAttentionAutoEncoder(resolution, num_slots, num_iterations, cfg=cfg).to(device)
num_params = sum(p.numel() for p in model.parameters())
print('Number of parameters:')
--
GitLab