def test_result_collection_changes_device():
    """Test that the keys in the ResultCollection are moved to the device together with the collection."""
    results = _ResultCollection(training=True)
    fx, name = "training_step", "step_log_val"
    log_val = torch.tensor(7.0, device="cuda:0")

    # same device as the original tensor
    results.log(fx, name, log_val, on_step=True, on_epoch=False, reduce_fx="mean")
    assert results[f"{fx}.{name}"].cumulated_batch_size.device == log_val.device

    # moved to cpu
    results.cpu()
    assert results[f"{fx}.{name}"].cumulated_batch_size.device == torch.device("cpu")

    # same device as the new tensor
    results.log(fx, name, log_val, on_step=True, on_epoch=False, reduce_fx="mean")
    assert results[f"{fx}.{name}"].cumulated_batch_size.device == log_val.device
Test that the keys in the ResultCollection are moved to the device together with the collection.
tests/tests_pytorch/trainer/logging_/test_logger_connector.py