shardTensor metadata recalc after checkpoint state_dict (pytorch#4146)

Summary:
Pull Request resolved: pytorch#4146

X-link: facebookresearch/FBGEMM#1227

Add ST metadata recalc into the DistributedCheckpointWrapper, so that all the state_dict calls will invoke recalc virtual PMT.

Reviewed By: pradeepfn

Differential Revision: D73567632

fbshipit-source-id: 40dbdaa6f51a0d58dcb48008688c1a9a3c8939e3

Author: duduyi2013

fbgemm_gpu/fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py

@@ -159,6 +159,8 @@ def is_pinned(self):
 
     @property
     def dtype(self) -> torch.dtype:
+        if isinstance(self._wrapped, torch.Tensor):
+            return self._wrapped.dtype
         mapping = {"c10::Half": "half"}
         dtype_str: str = self._wrapped.dtype_str
         dtype_str = mapping.get(dtype_str, dtype_str)
@@ -169,14 +171,17 @@ def dtype(self) -> torch.dtype:
 
     @property
     def device(self) -> torch.device:
+        if isinstance(self._wrapped, torch.Tensor):
+            return self._wrapped.device
         device_str: str = self._wrapped.device_str
         device = torch.device(device_str)
         assert isinstance(device, torch.device)
         return device
 
     @property
     def layout(self) -> torch.layout:
-        pass
+        if isinstance(self._wrapped, torch.Tensor):
+            return self._wrapped.layout
         layout_str_mapping = {
             "SparseCsr": "sparse_csr",
             "Strided": "strided",