From 2802499a75d190bdae0910ad7d35d02059e5da05 Mon Sep 17 00:00:00 2001 From: alexcbb <alexchapin@hotmail.fr> Date: Mon, 24 Jul 2023 17:04:25 +0200 Subject: [PATCH] Fix problem when unpacking data --- osrt/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/osrt/model.py b/osrt/model.py index a653271..460a967 100644 --- a/osrt/model.py +++ b/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. -- GitLab