minor change

geo_deep_learning/tasks_with_models/segmentation_segformer.py

@@ -15,7 +15,7 @@ def __init__(self,
                  class_labels: List[str] = None,
                  **kwargs: Any):
         super().__init__()
-        self.save_hyperparameters(ignore=["loss"])
+        self.save_hyperparameters()
         self.model = SegFormer(encoder, in_channels, num_classes)
         self.loss = loss
         self.metric= MulticlassJaccardIndex(num_classes=num_classes, average=None, zero_division=np.nan)
@@ -60,7 +60,7 @@ def test_step(self, batch, batch_idx):
         loss = self.loss(y_hat, y)
         y_hat = y_hat.argmax(dim=1)
         test_metrics = self.classwise_metric(y_hat, y)
-        test_metrics["test_loss"] = loss
+        test_metrics["loss"] = loss
         self.log_dict(test_metrics, 
                       prog_bar=True, logger=True, 
                       on_step=False, on_epoch=True, sync_dist=True, rank_zero_only=True)