This PR removes .to("cuda") calls that cause the metric to be forcefully placed on GPU 0, causing a memory imbalance. A torchmetrics Metric is a nn.Module, and therefore it will be automatically placed on the right device when used in a LightningModule together with the Trainer.
The best practice in Lightning is to not manipulate module or tensor placement unless there is a good reason for it.
This PR removes
.to("cuda")calls that cause the metric to be forcefully placed on GPU 0, causing a memory imbalance. AtorchmetricsMetric is ann.Module, and therefore it will be automatically placed on the right device when used in a LightningModule together with the Trainer.The best practice in Lightning is to not manipulate module or tensor placement unless there is a good reason for it.