fix(test): pytest error

open-mmlab · Jun 2, 2022 · ee56d66 · ee56d66
1 parent cf3a54b
commit ee56d66
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/tests/test_pytorch/test_pytorch_ops.py b/tests/test_pytorch/test_pytorch_ops.py
@@ -137,15 +137,27 @@ def test_normal(self):
  model = OpModel(torch.nn.functional.linear, w, bias).eval()
  nodes = get_model_onnx_nodes(model, x)
  print(nodes)
- assert nodes[0].op_type == 'MatMul'
+ gemm_exist = False
+ for node in nodes:
+ if node.op_type == 'Gemm':
+ gemm_exist = True
+ break
+
+ assert gemm_exist is True
 
  def test_no_bias(self):
  x = torch.rand(1, 2, 3)
  w = torch.rand(2, 3)
  model = OpModel(torch.nn.functional.linear, w).eval()
  nodes = get_model_onnx_nodes(model, x)
  print(nodes)
- assert nodes[0].op_type == 'MatMul'
+ gemm_exist = False
+ for node in nodes:
+ if node.op_type == 'Gemm':
+ gemm_exist = True
+ break
+
+ assert gemm_exist is True
 
 
 @pytest.mark.usefixtures('prepare_symbolics')