From 6a3a8fc5a5867f1c79ee9e0489ea48423a3cb8b4 Mon Sep 17 00:00:00 2001 From: Matt Li Date: Wed, 23 Oct 2024 17:39:52 -0400 Subject: [PATCH] solve `different device` issue by adding `.to(self.device)` at output I met this problem at evaluation stage with `ns-eval --load-config .....` --- fruit_nerf/fruit_nerf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fruit_nerf/fruit_nerf.py b/fruit_nerf/fruit_nerf.py index 004db82..3322071 100644 --- a/fruit_nerf/fruit_nerf.py +++ b/fruit_nerf/fruit_nerf.py @@ -404,7 +404,7 @@ def get_image_metrics_and_images( self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: image = batch["image"].to(self.device) - rgb = outputs["rgb"] + rgb = outputs["rgb"].to(self.device) rgb = torch.clamp(rgb, min=0, max=1) acc = colormaps.apply_colormap(outputs["accumulation"]) depth = colormaps.apply_depth_colormap( @@ -452,7 +452,7 @@ def get_image_metrics_and_images( from torchmetrics.classification import BinaryJaccardIndex metric = BinaryJaccardIndex().to(self.device) semantic_labels = torch.nn.functional.softmax(outputs["semantics"]) - iou = metric(semantic_labels[..., 0], batch["fruit_mask"][..., 0]) + iou = metric(semantic_labels[..., 0].to(self.device), batch["fruit_mask"][..., 0]) metrics_dict["iou"] = float(iou) return metrics_dict, images_dict