def test_to_device(accelerator, expected):
    """Test that the to_device method can move various objects to the device determined by the accelerator."""
    fabric = Fabric(accelerator=accelerator, devices=1)
    fabric.launch()

    expected_device = torch.device(expected)

    # module
    module = torch.nn.Linear(2, 3)
    module = fabric.to_device(module)
    assert all(param.device == expected_device for param in module.parameters())

    # tensor
    tensor = torch.rand(2, 2)
    tensor = fabric.to_device(tensor)
    assert tensor.device == expected_device

    # collection
    collection = {"data": torch.rand(2, 2), "int": 1}
    collection = fabric.to_device(collection)
    assert collection["data"].device == expected_device
Test that the to_device method can move various objects to the device determined by the accelerator.
tests/tests_fabric/test_fabric.py