fixed cpu backup issue

pytorch · Jul 9, 2023 · 5311353 · 5311353
1 parent 5d50a78
commit 5311353
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/aten/src/ATen/native/mps/operations/UnaryOps.mm b/aten/src/ATen/native/mps/operations/UnaryOps.mm
@@ -388,7 +388,12 @@ void cumulative_op_impl(const Tensor& self,
               ")");
   if (!is_macos_13_or_newer()) {
     TORCH_WARN_ONCE(op_name, " supported by MPS on MacOS 13+, please upgrade");
-    auto cpu_result = self.to(at::Device(kCPU)).cumsum(dim, dtype);
+    Tensor cpu_result;
+    if (cumulativeOpType == MPSCumulativeOpType::CUMSUM) {
+      cpu_result = self.to(at::Device(kCPU)).cumsum(dim, dtype);
+    } else if (cumulativeOpType == MPSCumulativeOpType::CUMPROD) {
+      cpu_result = self.to(at::Device(kCPU)).cumprod(dim, dtype);
+    }
     at::_copy_from_and_resize(cpu_result, result);
     return;
   }