Closed KubaMichalczyk closed 1 month ago
If you're working with multiple GPUs and the input tensor x is not on cuda:0, the following line causes the tensor to land on the default GPU. This causes a mismatch in any later call on x, as other tensors are likely still on the original GPU.
x
cuda:0
https://github.com/Lightning-AI/torchmetrics/blob/50e7da3de21fd50c39fe8ff07b99bba2ce4be5e6/src/torchmetrics/utilities/data.py#L218
A simple reference to x.device can solve this problem:
x.device
return x.cpu().cumsum(dim=dim, dtype=dtype).cuda(x.device)
Hi! thanks for your contribution!, great first issue!
If you're working with multiple GPUs and the input tensor
x
is not oncuda:0
, the following line causes the tensor to land on the default GPU. This causes a mismatch in any later call onx
, as other tensors are likely still on the original GPU.https://github.com/Lightning-AI/torchmetrics/blob/50e7da3de21fd50c39fe8ff07b99bba2ce4be5e6/src/torchmetrics/utilities/data.py#L218
A simple reference to
x.device
can solve this problem:return x.cpu().cumsum(dim=dim, dtype=dtype).cuda(x.device)