Resolve issue on output model

osrt/model.py

@@ -121,7 +121,7 @@ class LitSlotAttentionAutoEncoder(pl.LightningModule):
         masks = masks.softmax(dim = 1)
         recon_combined = (recons * masks).sum(dim = 1)
 
-        return recon_combined, recons, masks, slots, attn_slotwise.unsqueeze(-2).unflatten(-1, x.shape[-2:]) if attn_slotwise else None
+        return recon_combined, recons, masks, slots, attn_slotwise.unsqueeze(-2).unflatten(-1, x.shape[-2:]) if attn_slotwise is not None else None
 
     def training_step(self, batch, batch_idx):
         """Perform a single training step."""

train_sa.py

@@ -81,8 +81,6 @@ def main():
 
     trainer.fit(model, train_loader, val_loader)
 
-
-
     #### Create datasets
     vis_dataset = data.get_dataset('train', cfg['data'])
     vis_loader = DataLoader(

visualize_sa.py

@@ -52,7 +52,7 @@ def main():
     #### Create model
     model = LitSlotAttentionAutoEncoder(resolution, num_slots, num_iterations, cfg=cfg).to(device)
     checkpoint = torch.load(args.ckpt)
-    
+
     model.load_state_dict(checkpoint['state_dict'])
 
     model.eval()