Bug on spleen_segmentation_3d_visualization_basic.ipynb

[mingxin-zheng]

Describe the bug
Tensors on different device causes a bug when running the last code cell in 3d_segmentation/spleen_segmentation_3d_visualization_basic.ipynb

To Reproduce
Run the notebook and it will generate such error in MONAI 1.1

Error message in the last code cell

---

RuntimeError Traceback (most recent call last)
Cell In [17], line 15
13 val_outputs, val_labels = from_engine(["pred", "label"])(val_data)
14 # compute metric for current iteration
---> 15 dice_metric(y_pred=val_outputs, y=val_labels)
17 # aggregate the final mean dice result
18 metric_org = dice_metric.aggregate().item()

File /workspace/monai/monai-in-dev/monai/metrics/metric.py:330, in CumulativeIterationMetric.call(self, y_pred, y)
313 def call(self, y_pred: TensorOrList, y: Optional[TensorOrList] = None):
314 """
315 Execute basic computation for model prediction and ground truth.
316 It can support both list of channel-first Tensor and batch-first Tensor.
(...)
328 a batch-first tensor (BC[HWD]) or a list of batch-first tensors.
329 """
--> 330 ret = super().call(y_pred=y_pred, y=y)
331 if isinstance(ret, (tuple, list)):
332 self.extend(*ret)

File /workspace/monai/monai-in-dev/monai/metrics/metric.py:67, in IterationMetric.call(self, y_pred, y)
65 # handling a list of channel-first data
66 if isinstance(y_pred, (list, tuple)) or isinstance(y, (list, tuple)):
---> 67 return self._compute_list(y_pred, y)
68 # handling a single batch-first data
69 if isinstance(y_pred, torch.Tensor):

File /workspace/monai/monai-in-dev/monai/metrics/metric.py:89, in IterationMetric.compute_list(self, y_pred, y)
75 """
76 Execute the metric computation for y_pred and y in a list of "channel-first" tensors.
77
(...)
86 Note: subclass may enhance the operation to have multi-thread support.
87 """
88 if y is not None:
---> 89 ret = [self.compute_tensor(p.detach().unsqueeze(0), y.detach().unsqueeze(0)) for p, y in zip(y_pred, y)]
90 else:
91 ret = [self.compute_tensor(p.detach().unsqueeze(0), None) for p_ in y_pred]

File /workspace/monai/monai-in-dev/monai/metrics/metric.py:89, in (.0)
75 """
76 Execute the metric computation for y_pred and y in a list of "channel-first" tensors.
77
(...)
86 Note: subclass may enhance the operation to have multi-thread support.
87 """
88 if y is not None:
---> 89 ret = [self.compute_tensor(p.detach().unsqueeze(0), y.detach().unsqueeze(0)) for p, y_ in zip(y_pred, y)]
90 else:
91 ret = [self.compute_tensor(p.detach().unsqueeze(0), None) for p_ in y_pred]

File /workspace/monai/monai-in-dev/monai/metrics/meandice.py:83, in DiceMetric._compute_tensor(self, y_pred, y)
81 raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
82 # compute dice (BxC) for each channel for each batch
---> 83 return compute_dice(
84 y_pred=y_pred, y=y, include_background=self.include_background, ignore_empty=self.ignore_empty
85 )

File /workspace/monai/monai-in-dev/monai/metrics/meandice.py:143, in compute_dice(y_pred, y, include_background, ignore_empty)
141 n_len = len(y_pred.shape)
142 reduce_axis = list(range(2, n_len))
--> 143 intersection = torch.sum(y * y_pred, dim=reduce_axis)
145 y_o = torch.sum(y, reduce_axis)
146 y_pred_o = torch.sum(y_pred, dim=reduce_axis)

File /workspace/monai/monai-in-dev/monai/data/meta_tensor.py:268, in MetaTensor.torch_function(cls, func, types, args, kwargs)
266 if kwargs is None:
267 kwargs = {}
--> 268 ret = super().torch_function(func, types, args, kwargs)
269 # if out has been used as argument, metadata is not copied, nothing to do.
270 # if "out" in kwargs:
271 # return ret
272 if _not_requiring_metadata(ret):

File /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:1265, in Tensor.torch_function(cls, func, types, args, kwargs)
1262 return NotImplemented
1264 with _C.DisableTorchFunction():
-> 1265 ret = func(*args, **kwargs)
1266 if func in get_default_nowrap_functions():
1267 return ret

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Environment (please complete the following information):

• MONAI version: 1.1.0+2.g97918e46
• Numpy version: 1.22.2
• Pytorch version: 1.13.0a0+d0d6b1f
• MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False
• MONAI rev id: 97918e46e0d2700c050e678d72e3edb35afbd737
• MONAI file: /workspace/monai/monai-in-dev/monai/init.py

Additional context

for k, v in val_data[0].items():
    try:
        print(k, v.device)
    except:
        pass

image cpu
label cpu
foreground_start_coord cpu
foreground_end_coord cpu
pred cuda:0