[Bug][Relay] fix relay frontend pytorch op addmm bug

python/tvm/relay/frontend/pytorch.py

@@ -1549,24 +1549,22 @@ def flatten(self, inputs, input_types):
     def addmm(self, inputs, input_types):
         input_mat = inputs[0]
         mat1 = inputs[1]
-        data_type = input_types[1]
         mat2 = inputs[2]
-
         beta = inputs[3]
         alpha = inputs[4]
-
+        data_type = input_types[1]
+
+        transposed_mat2 = _op.transform.transpose(mat2, axes=[1, 0])
+        units = self.infer_shape(transposed_mat2)[0]
+        dense_out = _op.nn.dense(mat1, transposed_mat2, units=units)
+
         if not isinstance(alpha, _expr.Expr) and alpha != 1:
             alpha = _create_typed_const(alpha, data_type)
-            mat1 *= alpha
+            dense_out *= alpha
 
         if not isinstance(beta, _expr.Expr) and beta != 1:
             beta = _create_typed_const(beta, data_type)
-            mat2 *= beta
-
-        transposed_mat2 = _op.transform.transpose(mat2, axes=[1, 0])
-
-        units = self.infer_shape(transposed_mat2)[0]
-        dense_out = _op.nn.dense(mat1, transposed_mat2, units=units)
+            input_mat *= beta
 
         return dense_out + input_mat
 

tests/python/frontend/pytorch/test_forward.py

@@ -5260,6 +5260,20 @@ def test_weight_norm():
     verify_model(linear_wn.eval().float(), input_data_linear)
 
 
+@tvm.testing.uses_gpu
+def test_addmm():
+    def test_fn(alpha, beta):
+        return lambda inp, batch1, batch2: torch.addmm(
+            inp, batch1, batch2, beta=beta, alpha=alpha
+        )
+
+    M = torch.randn(3, 5)
+    batch1 = torch.randn(3, 4)
+    batch2 = torch.randn(4, 5)
+
+    verify_model(test_fn(0.4, 0.8), [M, batch1, batch2])
+
+
 @tvm.testing.uses_gpu
 def test_baddbmm():
     def test_fn(alpha, beta):