def test_fabric_dataloader_device_placement(src_device_str, dest_device_str):
    """Test that the FabricDataLoader moves data to the device in its iterator."""
    src_device = torch.device(src_device_str)
    dest_device = torch.device(dest_device_str)

    sample0 = torch.tensor(0, device=src_device)
    sample1 = torch.tensor(1, device=src_device)
    sample2 = {"data": torch.tensor(2, device=src_device)}
    sample3 = {"data": torch.tensor(3, device=src_device)}
    dataloader = DataLoader([sample0, sample1, sample2, sample3], batch_size=2)
    fabric_dataloader = _FabricDataLoader(dataloader=dataloader, device=dest_device)
    iterator = iter(fabric_dataloader)

    batch0 = next(iterator)
    assert torch.equal(batch0, torch.tensor([0, 1], device=dest_device))

    batch1 = next(iterator)
    assert torch.equal(batch1["data"], torch.tensor([2, 3], device=dest_device))
Test that the FabricDataLoader moves data to the device in its iterator.
tests/tests_fabric/test_wrappers.py