Merge pull request #120 from IGNF/119-fix-ddp-regression

Fix wrong device error created by a breaking change in torch when using DDP

CHANGELOG.md

@@ -1,4 +1,8 @@
 # CHANGELOG
+
+### 3.8.4
+- fix: move IoU appropriately to fix wrong device error created by a breaking change in torch when using DDP.
+
 ### 3.8.3
 - fix: prepare_data_per_node is a flag and was incorrectly used as a replacement for prepare_data.
 

myria3d/_version.py

@@ -1,4 +1,4 @@
-__version__ = "3.8.3"
+__version__ = "3.8.4"
 
 
 if __name__ == "__main__":

myria3d/models/model.py

@@ -148,7 +148,7 @@ def training_step(self, batch: Batch, batch_idx: int) -> dict:
         return {"loss": loss, "logits": logits, "targets": targets}
 
     def on_train_epoch_end(self) -> None:
-        iou_epoch = self.train_iou.compute()
+        iou_epoch = self.train_iou.to(self.device).compute()
         self.log("train/iou", iou_epoch, on_step=False, on_epoch=True, prog_bar=True)
         self.log_all_class_ious(self.train_iou.confmat, "train")
         log_comet_cm(self, self.train_iou.confmat, "train")
@@ -186,7 +186,7 @@ def on_validation_epoch_end(self) -> None:
             outputs : output of validation_step
 
         """
-        iou_epoch = self.val_iou.compute()
+        iou_epoch = self.val_iou.to(self.device).compute()
         self.log("val/iou", iou_epoch, on_step=False, on_epoch=True, prog_bar=True)
         self.log_all_class_ious(self.val_iou.confmat, "val")
         log_comet_cm(self, self.val_iou.confmat, "val")
@@ -221,7 +221,7 @@ def on_test_epoch_end(self) -> None:
             outputs : output of test
 
         """
-        iou_epoch = self.test_iou.compute()
+        iou_epoch = self.test_iou.to(self.device).compute()
         self.log("test/iou", iou_epoch, on_step=False, on_epoch=True, prog_bar=True)
         self.log_all_class_ious(self.test_iou.confmat, "test")
         log_comet_cm(self, self.test_iou.confmat, "test")