Add device property to lazy load functionality (#20183)

Author: t-vi

src/lightning/fabric/utilities/load.py

@@ -140,6 +140,10 @@ def __torch_function__(
         loaded_args = [(arg._load_tensor() if isinstance(arg, _NotYetLoadedTensor) else arg) for arg in args]
         return func(*loaded_args, **kwargs)
 
+    @property
+    def device(self) -> torch.device:
+        return torch.device(self.storageinfo[3])
+
     def __getattr__(self, name: str) -> Any:
         # These properties don't require materialization and can be accessed through the meta tensor directly
         if name in {
@@ -160,7 +164,7 @@ def __getattr__(self, name: str) -> Any:
             return getattr(self.metatensor, name)
 
         # materializing these is needed for quantization (see lit-gpt)
-        if name in {"contiguous", "cuda", "half", "data"}:
+        if name in {"contiguous", "cuda", "half", "data", "to"}:
             return getattr(self._load_tensor(), name)
 
         raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

tests/tests_fabric/utilities/test_load.py

@@ -31,6 +31,8 @@ def test_lazy_load_module(tmp_path):
     model1.load_state_dict(checkpoint)
 
     assert isinstance(checkpoint["weight"], _NotYetLoadedTensor)
+    assert checkpoint["weight"].device == torch.device("cpu")
+    assert type(checkpoint["weight"].to("cpu")) is torch.Tensor
     assert type(model0.weight.data) is torch.Tensor
     assert torch.equal(model0.weight, model1.weight)
     assert torch.equal(model0.bias, model1.bias)