Fix problem when unpacking data

osrt/model.py

@@ -159,7 +159,7 @@ class LitSlotAttentionAutoEncoder(pl.LightningModule):
 
         # Get the prediction of the model and compute the loss.
         preds = self.one_step(input_image)
-        recon_combined, recons, masks, slots = preds
+        recon_combined, recons, masks, slots, _ = preds
         input_image = input_image.permute(0, 2, 3, 1)
         loss_value = criterion(recon_combined, input_image)
         del recons, masks, slots  # Unused.
@@ -179,7 +179,7 @@ class LitSlotAttentionAutoEncoder(pl.LightningModule):
 
         # Get the prediction of the model and compute the loss.
         preds = self.one_step(input_image)
-        recon_combined, recons, masks, slots = preds
+        recon_combined, recons, masks, slots, _ = preds
         input_image = input_image.permute(0, 2, 3, 1)
         loss_value = criterion(recon_combined, input_image)
         del recons, masks, slots  # Unused.