with getattr(torch, "_C").DisableTorchFunctionSubclass():
data_descriptor.__set__(self, local_value)
data_descriptor.__set__(self._local_tensor, local_value)
if requires_grad is not None and hasattr(self, "requires_grad_"):
self.requires_grad_(requires_grad)
def cpu(self):
"""Move the local shard to CPU while preserving DTensor metadata."""
# pylint: disable=C0415
from hyper_parallel.core.dtensor.dtensor import DTensor
return DTensor.from_local(
self._local_tensor.cpu(),
self._device_mesh,
self._alias_placements(),
)