RMI Loss with densenet121

code/model.py

@@ -30,7 +30,7 @@ class TestModule(pl.LightningModule):
 
         # Define the model
         self.model = smp.Unet(
-            encoder_name="resnet34",  # Also consider using smaller or larger encoders
+            encoder_name="densenet121",  # Also consider using smaller or larger encoders
             encoder_weights= "imagenet",  # Do the pretrained weights help? Try with or without
             in_channels=1,  # We use 1 chanel transmittance as input
             classes=1,  # classes == output channels. We use one output channel for cyto data