def test_modules_without_parameters(distributed, tmp_path):
    """Test that TorchMetrics get moved to the device despite not having any parameters."""

    class MetricsModel(TensorParallelModel):
        def __init__(self):
            super().__init__()
            self.metric = Accuracy("multiclass", num_classes=10)
            assert self.metric.device == self.metric.tp.device == torch.device("cpu")

        def setup(self, stage) -> None:
            assert self.metric.device == self.metric.tp.device == torch.device("cpu")

        def training_step(self, batch):
            assert self.metric.device.type == self.metric.tp.device.type == "cuda"
            self.metric(torch.rand(2, 10, device=self.device), torch.randint(0, 10, size=(2,), device=self.device))
            return super().training_step(batch)

    model = MetricsModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cuda",
        devices=2,
        strategy=ModelParallelStrategy(),
        max_steps=1,
        enable_checkpointing=False,
        logger=False,
    )
    trainer.fit(model)
Test that TorchMetrics get moved to the device despite not having any parameters.
tests/tests_pytorch/strategies/test_model_parallel_integration.py