From ebd32dc92027837df1157cb79436149d89d1fbd9 Mon Sep 17 00:00:00 2001 From: valhassan Date: Sat, 28 Sep 2024 10:06:40 -0400 Subject: [PATCH] minor change --- geo_deep_learning/tasks_with_models/segmentation_segformer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/geo_deep_learning/tasks_with_models/segmentation_segformer.py b/geo_deep_learning/tasks_with_models/segmentation_segformer.py index b5db76dd..1b14555b 100644 --- a/geo_deep_learning/tasks_with_models/segmentation_segformer.py +++ b/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)