diff --git a/src/huggingface_hub/serialization/_torch.py b/src/huggingface_hub/serialization/_torch.py index 3c8b2b33f3..e87e728d23 100644 --- a/src/huggingface_hub/serialization/_torch.py +++ b/src/huggingface_hub/serialization/_torch.py @@ -368,18 +368,21 @@ def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]: return unique_id -def get_torch_storage_id(tensor: "torch.Tensor") -> Tuple["torch.device", Union[int, Tuple[Any, ...]], int]: +def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[Tuple["torch.device", Union[int, Tuple[Any, ...]], int]]: """ Return unique identifier to a tensor storage. - Multiple different tensors can share the same underlying storage. For - example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is + Multiple different tensors can share the same underlying storage. This identifier is guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with non-overlapping lifetimes may have the same id. + In the case of meta tensors, we return None since we can't tell if they share the same storage. Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278. """ - return tensor.device, _get_unique_id(tensor), get_torch_storage_size(tensor) + if tensor.device.type == "meta": + return None + else: + return tensor.device, _get_unique_id(tensor), get_torch_storage_size(tensor) def get_torch_storage_size(tensor: "torch.Tensor") -> int: