From c6bc5db90ba17e6c3c2e38cfceb1e0ef681eddf6 Mon Sep 17 00:00:00 2001 From: Ashwin Hari Date: Wed, 20 Mar 2024 11:50:33 -0700 Subject: [PATCH] add case for ort --- torch/csrc/Storage.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torch/csrc/Storage.cpp b/torch/csrc/Storage.cpp index 93dbc9c09bb2b..c22e6f5d1b95d 100644 --- a/torch/csrc/Storage.cpp +++ b/torch/csrc/Storage.cpp @@ -355,6 +355,8 @@ static PyObject* THPStorage_pynew( } else if (device.type() == at::DeviceType::PrivateUse1) { at::globalContext().lazyInitPrivateUse1(); allocator = c10::GetAllocator(device.type()); + } else if (device.type() == at::DeviceType::ORT) { + allocator = c10::GetAllocator(device.type()); } else { // NOLINTEND(bugprone-branch-clone) TORCH_CHECK(