Skip to content

Commit 9458b83

Browse files
Nitin-x2apytorchmergebot
authored andcommitted
[HPU] Add HPU as a supported device for NestedTensor (pytorch#148659)
This change enables basic NestedTensor operations on HPU, fixing the runtime error when creating a NestedTensor on HPU. - Extended `NestedTensorImpl` to recognize `hpu` as a valid storage device. - Added `NestedTensorHPU` to `DispatchKey` parsing in `DispatchKey.cpp`. - Updated `torchgen/model.py` to include `NestedTensorHPU` in `dispatch_keys`. - Modified `native_functions.yaml` to enable `NestedTensorHPU` support for various ops. Fixes #ISSUE_NUMBER Pull Request resolved: pytorch#148659 Approved by: https://github.com/jeromean, https://github.com/albanD, https://github.com/sujoysaraswati
1 parent 9aca001 commit 9458b83

File tree

4 files changed

+91
-89
lines changed

4 files changed

+91
-89
lines changed

aten/src/ATen/NestedTensorImpl.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ NestedTensorImpl::NestedTensorImpl(
182182
"coverage, and works with torch.compile.");
183183
auto storage_device = storage_.device();
184184
TORCH_INTERNAL_ASSERT(
185-
storage_device.is_cpu() || storage_device.is_cuda() || storage_device.is_xpu() || storage_device.is_privateuseone(),
185+
storage_device.is_cpu() || storage_device.is_cuda() || storage_device.is_xpu() || storage_device.is_hpu() || storage_device.is_privateuseone(),
186186
"NestedTensorImpl storage must be either CUDA, CPU, XPU or ", get_privateuse1_backend(), " but got ",
187187
storage_device);
188188
validate_nested_tensor_metadata(nested_sizes_, nested_strides_, storage_offsets_);

0 commit comments

Comments
 (0)