Fix resize matrix_power.out dynamic shapes (pytorch#113695)

Fixes pytorch#113003 Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#113695 Approved by: https://github.com/bdhirsh, https://github.com/lezcano
sands-lab · Nov 16, 2023 · c4c45ab · c4c45ab
1 parent 8a183bf
commit c4c45ab
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp
@@ -626,7 +626,7 @@ Tensor linalg_matrix_power_impl(
   if (_out.has_value()) {
     checkSameDevice("matrix_power", out, self);
     checkLinalgCompatibleDtype("matrix_power", out, self);
-    at::native::resize_output(out, self.sizes());
+    at::native::resize_output_symint(out, self.sym_sizes());
   }
 
   // For n=0 we return the identity matrix of the same shape as input.

diff --git a/aten/src/ATen/native/Resize.cpp b/aten/src/ATen/native/Resize.cpp
@@ -184,6 +184,11 @@ static void _maybe_resize_storage(TensorImpl* self, int64_t new_size_bytes) {
 }
 
 static void _maybe_resize_storage(TensorImpl* self, c10::SymInt new_size_bytes) {
+  if (self->is_cpu()) {
+    maybe_resize_storage_cpu(self, new_size_bytes.expect_int());
+    return;
+  }
+  TORCH_INTERNAL_ASSERT(self->is_meta());
   maybe_resize_storage_meta(self, std::move(new_size_bytes));
 }