Change MSE loss

osrt/layers.py

@@ -295,7 +295,6 @@ class SoftPositionEmbed(nn.Module):
     def forward(self, inputs):
         return inputs + self.dense(torch.tensor(self.grid).cuda()).permute(0, 3, 1, 2) #  from [b, h, w, c] to [b, c, h, w]
     
-    
 ### New transformer implementation of SlotAttention inspired from https://github.com/ThomasMrY/VCT/blob/master/models/visual_concept_tokenizor.py
 class TransformerSlotAttention(nn.Module):
     """

train_sa.py

@@ -43,7 +43,7 @@ def eval_step(batch, model, device):
     preds = model(input_image)
     recon_combined, recons, masks, slots = preds
     input_image = input_image.permute(0, 2, 3, 1)
-    loss_value = nn.MSELoss(recon_combined, input_image)
+    loss_value = F.mse_loss(recon_combined, input_image)
     del recons, masks, slots  # Unused.
     psnr = mse2psnr(loss_value)