MPS: Fix handling of 1D tensors in linear backward (pytorch#80759)

Fixes #pytorch#79784 Pull Request resolved: pytorch#80759 Approved by: https://github.com/ezyang
kulinseth · Jul 9, 2022 · 87dc222 · 87dc222
1 parent 8b7ac61
commit 87dc222
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/aten/src/ATen/native/mps/operations/Linear.mm b/aten/src/ATen/native/mps/operations/Linear.mm
@@ -242,9 +242,9 @@ Tensor _mps_linear_backward_input(
     MPSGraphTensor *biasTensor_ = nil;
   };
 
-  auto grad_output_reshaped = grad_output.dim() > 2 ?
+  auto grad_output_reshaped = grad_output.dim() != 2 ?
     grad_output.reshape({-1, grad_output.size(grad_output.dim() - 1)}) : grad_output;
-  auto input_reshaped = input.dim() > 2 ? input.reshape({-1, input.size(input.dim() - 1)}) : input;
+  auto input_reshaped = input.dim() != 2 ? input.reshape({-1, input.size(input.dim() - 1)}) : input;
 
   TORCH_CHECK(grad_output_reshaped.is_mps());
   TORCH_CHECK(input_reshaped.is_mps());

diff --git a/test/test_mps.py b/test/test_mps.py
@@ -367,6 +367,12 @@ def _linear_helper(self, in_features, out_features, shape, bias=True, backward_p
                 self.assertEqual(cpu_linear.bias.grad.size(), mps_linear.bias.grad.size())
                 self.assertEqual(cpu_linear.bias.grad, mps_linear.bias.grad.to("cpu"), atol=8e-04, rtol=10.4e-05)
 
+    def test_linear1D(self):
+        self._linear_helper(in_features=2, out_features=3, shape=([2]), bias=True, backward_pass=False)
+
+    def test_linear1D_backward(self):
+        self._linear_helper(in_features=2, out_features=3, shape=([2]), bias=True, backward_pass=True)
+
     def test_linear2D(self):
         self._linear_helper(in_features=2, out_features=3, shape=((4, 2)), bias=True, backward_pass=False)