Merge pull request #570 from asi1024/fix-test-torch-112

Small test fixes for torch 1.12
pfnet · Jul 11, 2022 · c40cd02 · c40cd02
2 parents 4095ed2 + 9b78950
commit c40cd02
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/tests/pytorch_pfn_extras_tests/onnx_tests/test_annotate.py b/tests/pytorch_pfn_extras_tests/onnx_tests/test_annotate.py
@@ -374,8 +374,8 @@ def forward(self, *xs):
     anchor_node, pre_node, next_node = named_nodes['Anchor_0_start']
     # anchor_attrs = [a.name for a in anchor_node.attribute]
     assert pre_node is None
-    assert next_node.name == 'Concat_4'
+    assert next_node.name.startswith('Concat_')
     anchor_node, pre_node, next_node = named_nodes['Anchor_0_end']
     # anchor_attrs = [a.name for a in anchor_node.attribute]
-    assert pre_node.name == 'Split_10'
+    assert pre_node.name.startswith('Split_')
     assert next_node is None
diff --git a/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py b/tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py
@@ -2,13 +2,20 @@
 import torch
 import torchvision
 
+import pytorch_pfn_extras
 from pytorch_pfn_extras_tests.onnx_tests.utils import run_model_test
 
 
+if pytorch_pfn_extras.requires("1.12.0"):
+    resnet18_kwargs = {'weights': None}
+else:
+    resnet18_kwargs = {'pretrained': True}
+
+
 def test_eval_resnet18():
     torch.manual_seed(100)
     run_model_test(
-        torchvision.models.resnet.resnet18(pretrained=True),
+        torchvision.models.resnet.resnet18(**resnet18_kwargs),
         (torch.rand(1, 3, 224, 224),),
         rtol=1e-03,
         use_gpu=True,
@@ -19,7 +26,7 @@ def test_eval_resnet18():
 @pytest.mark.xfail
 def test_train_resnet18():
     run_model_test(
-        torchvision.models.resnet.resnet18(pretrained=True),
+        torchvision.models.resnet.resnet18(**resnet18_kwargs),
         (torch.rand(1, 3, 224, 224),),
         rtol=1e-03,
         use_gpu=True,