Simplify test

Author: Turakar

linear_operator/test/linear_operator_test_case.py

@@ -58,19 +58,6 @@ def _test_matmul(self, rhs):
         actual = evaluated.matmul(rhs)
         self.assertAllClose(to_dense(res), actual)
 
-    def _test_t_matmul(self, rhs):
-        with torch.no_grad():
-            linear_op = self.create_linear_op()
-            linear_op_copy = torch.clone(linear_op)
-            evaluated = self.evaluate_linear_op(linear_op_copy)
-            rhs_evaluated = to_dense(rhs)
-
-            # Test operator
-            res = linear_op._t_matmul(rhs)
-            actual = evaluated.mT.matmul(rhs_evaluated)
-            res_evaluated = to_dense(res)
-            self.assertAllClose(res_evaluated, actual)
-
     def _test_rmatmul(self, lhs):
         linear_op = self.create_linear_op().detach().requires_grad_(True)
         linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
@@ -405,9 +392,18 @@ def test_matmul_matrix(self):
         return self._test_matmul(rhs)
 
     def test_t_matmul_matrix(self):
-        linear_op = self.create_linear_op()
-        rhs = torch.randn(*linear_op.batch_shape, linear_op.size(-2), 4)
-        return self._test_t_matmul(rhs)
+        with torch.no_grad():
+            linear_op = self.create_linear_op()
+            rhs = torch.randn(*linear_op.batch_shape, linear_op.size(-2), 4)
+            linear_op_copy = torch.clone(linear_op)
+            evaluated = self.evaluate_linear_op(linear_op_copy)
+            rhs_evaluated = to_dense(rhs)
+
+            # Test operator
+            res = linear_op._t_matmul(rhs)
+            actual = evaluated.mT.matmul(rhs_evaluated)
+            res_evaluated = to_dense(res)
+            self.assertAllClose(res_evaluated, actual)
 
     def test_rmatmul_matrix(self):
         linear_op = self.create_linear_op()