Support resize on meta storage

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 46bf0449a95dbbf32a8b8a82c99bc6b118b448a2 Pull Request resolved: #101988
pytorch · May 22, 2023 · 3554daa · 3554daa
1 parent 38a2932
commit 3554daa
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/aten/src/ATen/native/Resize.h b/aten/src/ATen/native/Resize.h
@@ -34,6 +34,7 @@ TORCH_API bool resize_output_check(const Tensor& output, IntArrayRef shape);
 TORCH_API bool resize_output_check_symint(const Tensor& output, SymIntArrayRef shape);
 
 TORCH_API void resize_bytes_cpu(StorageImpl* storage, size_t size_bytes);
+TORCH_API void resize_bytes_meta(StorageImpl* storage, c10::SymInt size_bytes);
 
 static inline void maybe_resize_storage_cpu(TensorImpl* self, size_t new_size_bytes) {
   // It does not make sense to try to resize a storage

diff --git a/torch/csrc/StorageMethods.cpp b/torch/csrc/StorageMethods.cpp
@@ -123,6 +123,8 @@ static PyObject* THPStorage_resize_(PyObject* self, PyObject* number_arg) {
     const auto size_bytes = static_cast<size_t>(size_bytes_i);
     at::native::resize_bytes_cuda(storage.unsafeGetStorageImpl(), size_bytes);
 #endif
+  } else if (device_type == at::kMeta) {
+    at::native::resize_bytes_meta(storage.unsafeGetStorageImpl(), newsize);
   } else if (device_type == at::kPrivateUse1) {
     ptrdiff_t size_bytes_i = newsize;
     TORCH_CHECK(