Skip to content

Commit 2ea0970

Browse files
heidongxianhuapytorchmergebot
authored andcommitted
fix device type bug for custom device (pytorch#97213)
Fixes #ISSUE_NUMBER support the custom renamed device ,@bdhirsh , please review my changes. Pull Request resolved: pytorch#97213 Approved by: https://github.com/bdhirsh, https://github.com/kit1980
1 parent fcc312e commit 2ea0970

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

torch/_tensor.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def __deepcopy__(self, memo):
100100
or self.device.type in ["lazy", "xla", "mps", "ort", "meta", "ipu"]
101101
or (
102102
not torch._C._has_storage(self)
103-
and self.device.type == "privateuseone"
103+
and self.device.type == torch._C._get_privateuse1_backend_name()
104104
)
105105
or (type(self) is not Tensor and self.data_ptr() == 0)
106106
):
@@ -256,7 +256,8 @@ def _reduce_ex_internal(self, proto):
256256
# `tolist()` converts every single element in the tensor into python objects
257257
# and serialize them one by one.
258258
if self.device.type in ["xla", "ort"] or (
259-
not torch._C._has_storage(self) and self.device.type == "privateuseone"
259+
not torch._C._has_storage(self)
260+
and self.device.type == torch._C._get_privateuse1_backend_name()
260261
):
261262
# Convert BFloat16 tesors to Float32 before conversion to numpy, as numpy doesn't
262263
# support BFloat16. The rebuild tensor from numpy takes in the original self.dtype,

0 commit comments

Comments
 (0)