diff --git a/aten/src/ATen/TensorIterator.cpp b/aten/src/ATen/TensorIterator.cpp index dccd2009046..7b1442db75a 100644 --- a/aten/src/ATen/TensorIterator.cpp +++ b/aten/src/ATen/TensorIterator.cpp @@ -1491,10 +1491,19 @@ void TensorIteratorBase::build(TensorIteratorConfig& config) { if (is_meta_) return; + auto has_storage = true; + for (auto& op : operands_) { + has_storage &= op.tensor_base().has_storage(); + } + auto privateuse1_without_storage = + common_device_.type() == DeviceType::PrivateUse1 && + !has_storage; + // XLA and lazy tensors don't have storage, so they don't have an underlying data pointer. // Nothing beyond this point is important for meta functions, so it's fine to exit early here. // Extend the condition to ORT tesnors as ORT tensors also don't have storage. - if (common_device_.type() == DeviceType::XLA || + if (privateuse1_without_storage || + common_device_.type() == DeviceType::XLA || common_device_.type() == DeviceType::IPU || common_device_.type() == DeviceType::Lazy || common_device_.type() == DeviceType::ORT || diff --git a/torch/_tensor.py b/torch/_tensor.py index 05171c7ae0e..c7cb9b338b9 100644 --- a/torch/_tensor.py +++ b/torch/_tensor.py @@ -116,6 +116,7 @@ class Tensor(torch._C._TensorBase): if ( self.is_sparse or self.device.type in ["lazy", "xla", "mps", "ort", "meta", "hpu"] + or (self.storage is None and self.device.type == "privateuseone") or (type(self) is not Tensor and self.data_ptr() == 0) ): new_tensor = self.clone() @@ -271,7 +272,9 @@ class Tensor(torch._C._TensorBase): # 2. Python list is not a good fit due to performance reason. # `tolist()` converts every single element in the tensor into python objects # and serialize them one by one. - if self.device.type in ["xla", "ort", "hpu"]: + if self.device.type in ["xla", "ort", "hpu"] or ( + self.storage is None and self.device.type == "privateuseone" + ): # Convert BFloat16 tesors to Float32 before conversion to numpy, as numpy doesn't # support BFloat16. The rebuild tensor from numpy takes in the original self.dtype, # this would reconstruct the BFloat16 tensor from numpy.