def test_fabric_module_device_dtype_propagation(device_str, dtype):
    """Test that the FabricModule propagates device and dtype properties to its submodules (e.g. torchmetrics)."""
    device = torch.device(device_str)

    class DeviceModule(_DeviceDtypeModuleMixin):
        pass

    device_module = DeviceModule()
    fabric_module = _FabricModule(device_module, Mock())
    fabric_module.to(device)
    assert device_module.device == device
    assert fabric_module.device == device

    fabric_module.to(dtype)
    assert device_module.dtype == dtype
    assert fabric_module.dtype == dtype
Test that the FabricModule propagates device and dtype properties to its submodules (e.g. torchmetrics).
tests/tests_fabric/test_wrappers.py