Add check for backend of arguments to bmm cpu (#12434)

t-vi · facebook-github-bot · commit 0e44db8b0d23 · 2018-10-07T18:55:42.000-07:00
Summary: Fixes: #12406 Thank you, jcjohnson, for reporting. Pull Request resolved: #12434 Differential Revision: D10235799 Pulled By: soumith fbshipit-source-id: 44ee35010bac3791901f604095f5b4bc66b0e7f8
diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp
@@ -283,6 +283,7 @@ static inline Tensor& bmm_out_or_baddbmm_(Tensor& self_or_result, const Tensor&
   TensorArg self_arg(self_or_result, is_bmm_out ? "self" : "result", 0);
   TensorArg b1_arg(batch1, "batch1", 1);
   TensorArg b2_arg(batch2, "batch2", 2);
+  checkBackend(c, {self_or_result, batch1, batch2}, Backend::CPU);
   checkDim(c, b1_arg, 3);
   checkDim(c, b2_arg, 3);
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -1600,6 +1600,10 @@ def test_bmm(self):
         for i in range(num_batches):
             r = torch.mm(b1[i], b2[i])
             self.assertEqual(r, res[i])
+        if torch.cuda.is_available():
+            # check that mixed arguments are rejected
+            self.assertRaises(RuntimeError, lambda: torch.bmm(b1, b2.cuda()))
+            self.assertRaises(RuntimeError, lambda: torch.bmm(b1.cuda(), b2))
 
     def test_addbmm(self):
         # num_batches = 10
