commit a076a74f1118da171cf70d00d1de4abbe27cf85d
author Chen, Zejun <zejun.chen@intel.com> Wed Nov 29 05:59:31 2023 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Wed Nov 29 05:59:35 2023 +0000
tree 4877f0baf6081328741db1119c93ea653659bb32
parent 69c4819f530023dccda048ce1b9057b04a5f972c

[Nested Tensor] Add xpu device in assertion for nested tensor creation (#114664)

Add xpu device checking in nested tensor creation
Pull Request resolved: https://github.com/pytorch/pytorch/pull/114664
Approved by: https://github.com/jgong5, https://github.com/xunnanxu

aten/src/ATen/NestedTensorImpl.cpp

diff --git a/aten/src/ATen/NestedTensorImpl.cpp b/aten/src/ATen/NestedTensorImpl.cpp
index 99da6fb..223a95e 100644
--- a/aten/src/ATen/NestedTensorImpl.cpp
+++ b/aten/src/ATen/NestedTensorImpl.cpp
@@ -179,8 +179,8 @@
       "in the near future.");
   auto storage_device = storage_.device();
   TORCH_INTERNAL_ASSERT(
-      storage_device.is_cpu() || storage_device.is_cuda() || storage_device.is_privateuseone(),
-      "NestedTensorImpl storage must be either CUDA, CPU or ", get_privateuse1_backend(), " but got ",
+      storage_device.is_cpu() || storage_device.is_cuda() || storage_device.is_xpu() || storage_device.is_privateuseone(),
+      "NestedTensorImpl storage must be either CUDA, CPU, XPU or ", get_privateuse1_backend(), " but got ",
       storage_device);
   validate_nested_tensor_metadata(nested_sizes_, nested_strides_, storage_offsets_);
   refresh_dim();