def test_lazy_load_tensor(tmp_path):
    """Test that lazy load can handle different classes of tensors."""
    expected = {
        "tensor": torch.rand(2),
        "parameter": nn.Parameter(torch.rand(3)),
        "subclass": torch.Tensor._make_subclass(ATensor, torch.rand(4)),
    }
    torch.save(expected, tmp_path / "data.pt")

    loaded = _lazy_load(tmp_path / "data.pt")
    for t0, t1 in zip(expected.values(), loaded.values()):
        assert isinstance(t1, _NotYetLoadedTensor)
        t1_materialized = _materialize_tensors(t1)
        assert type(t0) == type(t1_materialized)  # noqa: E721
        assert torch.equal(t0, t1_materialized)
Test that lazy load can handle different classes of tensors.
tests/tests_fabric/utilities/test_load.py