From 49e6934873518b1ced0e9ca7498aa392035579b0 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sat, 26 Jul 2025 12:54:03 -0400 Subject: [PATCH 1/2] use untyped storage for dtensors due to deprecation --- src/transformers/pytorch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/pytorch_utils.py b/src/transformers/pytorch_utils.py index c3cc4579e5c6..b340254c4a18 100644 --- a/src/transformers/pytorch_utils.py +++ b/src/transformers/pytorch_utils.py @@ -299,7 +299,7 @@ def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]: if isinstance(tensor, DTensor): local_tensor = tensor.to_local() - return tensor.device, local_tensor.storage().data_ptr(), tensor.nbytes + return tensor.device, local_tensor.untyped_storage().data_ptr(), tensor.nbytes if tensor.device.type == "xla" and is_torch_xla_available(): # NOTE: xla tensors dont have storage From 16341a724df123be31c9385e77c353b3d6417104 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sat, 26 Jul 2025 13:08:22 -0400 Subject: [PATCH 2/2] use nbytes from storage --- src/transformers/pytorch_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/pytorch_utils.py b/src/transformers/pytorch_utils.py index b340254c4a18..bed115e72d3d 100644 --- a/src/transformers/pytorch_utils.py +++ b/src/transformers/pytorch_utils.py @@ -299,7 +299,7 @@ def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]: if isinstance(tensor, DTensor): local_tensor = tensor.to_local() - return tensor.device, local_tensor.untyped_storage().data_ptr(), tensor.nbytes + return tensor.device, local_tensor.untyped_storage().data_ptr(), tensor.untyped_storage().nbytes() if tensor.device.type == "xla" and is_torch_xla_available(): # NOTE: xla tensors dont have storage