code: 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)
          assert torch.equal(t0, t1_materialized)
  
file_path: pytorch-lightning/tests/tests_fabric/utilities/test_load.py
imports: import pytest
  import torch
  import torch.nn
  from lightning.fabric.utilities.load import _lazy_load, _materialize_tensors, _move_state_into, _NotYetLoadedTensor
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: test_lazy_load_tensor(tmp_path) without return types