Error when calling `.to('cpu')` on a CUDA metric loaded on a CPU-only machine

[tringwald]

🐛 Bug

torchmetrics seems unable to handle a metric that has been serialized on a CUDA-enabled installation when later loading on a CPU-only installation and calling metric.to('cpu') (see below). This even happens with torch.load(..., map_location='cpu').

# On CUDA-enabled torch
import torch
import torchmetrics 
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5).cuda()
print(metric._device)
# device(type='cuda', index=0)
torch.save(metric, 'test.pth')
# On CPU-only torch
m = torch.load('test.pth', map_location='cpu')
print(m._device)
# device(type='cuda', index=0)
m.to('cpu')
#[...]
#    raise AssertionError("Torch not compiled with CUDA enabled")
#AssertionError: Torch not compiled with CUDA enabled

The problematic code is here:

torchmetrics/src/torchmetrics/metric.py

Lines 811 to 813 in 894de4c

	# make sure to update the device attribute
	# if the dummy tensor moves device by fn function we should also update the attribute
	self._device = fn(torch.zeros(1, device=self.device)).device

After loading, self.device will still refer to the original CUDA device as that is what was serialized:

>>> m = torch.load('test.pth', map_location='cpu')
>>> m._device
device(type='cuda', index=0)

Constructing a tensor with torch.zeros(1, device=self.device) will then error out if CUDA is not available.

To Reproduce

see above

Expected behavior

Calling metric.to('cpu') on a metric stored on the CPU device should not throw an error.

Environment

torchmetrics==1.2.0

Additional context

Originally reported here:

• Cannot load a GPU-trained nn.Module on a CPU-only machine pytorch/pytorch#113973
• Add support to saving.py for loading GPU-trained models on CPU-only machines pytorch-lightning#19024

[github-actions]

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

[SkafteNicki]

Hi @tringwald, thanks for raising this issue.
The correct way to save and load metrics are saving/loading the state_dict instead of the hole metric, as any other nn.Module. The only difference is that you need to call the metric.persistent(True) method before. Here is an example.

import torch
from torchmetrics.classification import MulticlassAccuracy

metric = MulticlassAccuracy(num_classes=5).to("cuda")
metric.persistent(True)
metric.update(torch.randint(5, (100,)).cuda(), torch.randint(5, (100,)).cuda())
torch.save(metric.state_dict(), "metric.pth")

metric2 = MulticlassAccuracy(num_classes=5).to("cpu")
metric2.load_state_dict(torch.load("metric.pth", map_location="cpu"))

# These will match, but be on different devices
print(metric.metric_state)
print(metric2.metric_state)

This is however not clear at all from our documentation which will be fixed in PR #2237.

[amorehead]

@SkafteNicki, thank you for adding this new documentation on saving and loading metrics. However, I am entirely clear on what this would look like for a typical LightningModule. For example, I typically instantiate all my torchmetrics in my LightningModule's __init__() function. Is this best practice, or is there another place where I should be instantiating these? Moreover, how can I ensure that Lightning only saves my metrics' state_dict in a persistent way?