Debugging Validation and testing

lib/models/homotopy_classification/cnn_based.py

@@ -57,7 +57,8 @@ class ConvHomDetector(LightningBaseModule):
         # Model Parameters
         self.in_shape = self.dataset.map_shapes_max
         assert len(self.in_shape) == 3, f'Image or map shape has to have 3 dims, but had: {len(self.in_shape)}'
-        self.criterion = nn.BCEWithLogitsLoss()
+        self.criterion = nn.BCELoss()
+        self.sigmoid = nn.Sigmoid()
 
         # NN Nodes
         # ============================
@@ -100,4 +101,5 @@ class ConvHomDetector(LightningBaseModule):
         tensor = self.flatten(tensor)
         tensor = self.linear(tensor)
         tensor = self.classifier(tensor)
+        tensor = self.sigmoid(tensor)
         return tensor