Metrics fails or wrong output device with CUDA

[matteo-bastico]

Describe the bug
Hello, some metrics have bugs when used with CUDA tensors. In particular, Generalized Dice Score fails with the following error when both y_pred and y are on the same CUDA device.

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

Instead, in Average Surface Distance the device of the output is always cpu even if both outputs are on gpu.

To Reproduce
For Generalized Dice Score:

x = torch.randint(0, 2, (1, 6, 96, 96, 96)).to("cuda:0")
y = torch.randint(0, 2, (1, 6, 96, 96, 96)).to("cuda:0")
monai.metrics.compute_generalized_dice(x, y)

will raise the RuntimeError.

For Average Surface Distance

x = torch.randint(0, 2, (1, 6, 96, 96, 96)).to("cuda:0")
y = torch.randint(0, 2, (1, 6, 96, 96, 96)).to("cuda:0")
monai.metrics.compute_average_surface_distance(x, y)

will return a tensor([[...]], dtype=torch.float64) instead of tensor([[...]], device='cuda:0') (see Mean Dice for correct behavior).

Expected behavior
The metrics are computed without raising errors when inputs are on gpu and the returned output is on the same device.

Environment

Monai: 1.1.0

[Nic-Ma]

Hi @yiheng-wang-nv ,

Seems there is some issue here? Could you please help take a look and may submit a PR to fix it?

Thanks in advance.

[yiheng-wang-nv]

Hi @matteo-bastico , for Generalized Dice Score, I did not meet same issue. Could you please run the following code and post the exact monai version?

from monai.config import print_config

print_config()

for Average Surface Distance, let me fix the issue first.

[matteo-bastico]

HI @yiheng-wang-nv, the output of print_config is the following:

MONAI version: 1.1.0
Numpy version: 1.24.1
Pytorch version: 1.13.1+cu117
MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False
MONAI rev id: a2ec3752f54bfc3b40e7952234fbeb5452ed63e3
MONAI __file__: /home/mbastico/anaconda3/envs/MI-Seg/lib/python3.10/site-packages/monai/__init__.py
Optional dependencies:
Pytorch Ignite version: 0.4.10
Nibabel version: 5.0.0
scikit-image version: NOT INSTALLED or UNKNOWN VERSION.
Pillow version: 9.4.0
Tensorboard version: 2.11.2
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.14.1+cu117
tqdm version: 4.64.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: 5.9.4
pandas version: 1.5.2
einops version: 0.6.0
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.

[yiheng-wang-nv]

Hi @matteo-bastico , thanks for posting the versions. I double checked the code, it seems the latest pytorch dev version (1.14.0a0+410ce96) resolves the device issue on the following case (x is in cuda, but torch.tensor(1) is in cpu):

x = torch.randint(0, 2, (1, 6, 96, 96, 96)).to("cuda:0")
x[0] = torch.where(x[0] > 0, torch.tensor(1), torch.tensor(2))

Therefore, I did not not meet the same error. When I downgraded to 1.13.1, same error happens.
Let me fix it with another PR. Thanks again!