search

Lightning-AI / torchmetrics

Torchmetrics - Machine learning metrics for distributed, scalable PyTorch applications.
https://lightning.ai/docs/torchmetrics/
Apache License 2.0
2.07k stars 395 forks source link

_cumsum moves the data to default GPU when deterministic flag is on #2608

Closed KubaMichalczyk closed 1 month ago

KubaMichalczyk commented 2 months 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.

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)

github-actions[bot] commented 2 months ago

Hi! thanks for your contribution!, great first issue!