Update export test

Author: qubvel

tests/models/base.py

@@ -252,10 +252,11 @@ def test_compile(self):
             compiled_model(sample)
 
     @pytest.mark.torch_export
-    def test_torch_export(self):
+    def test_torch_export(self, eps=1e-5):
         if not check_run_test_on_diff_or_main(self.files_for_diff):
             self.skipTest("No diff and not on `main`.")
 
+        torch.manual_seed(42)
         sample = self._get_sample().to(default_device)
         model = self.get_default_model()
         model.eval()
@@ -271,7 +272,7 @@ def test_torch_export(self):
             exported_output = exported_model.module().forward(sample)
 
         self.assertEqual(eager_output.shape, exported_output.shape)
-        torch.testing.assert_close(eager_output, exported_output)
+        torch.testing.assert_close(eager_output, exported_output, rtol=eps, atol=eps)
 
     @pytest.mark.torch_script
     def test_torch_script(self):

tests/models/test_upernet.py

@@ -1,3 +1,5 @@
+import pytest
+
 from tests.models import base
 
 
@@ -6,3 +8,7 @@ class TestUnetModel(base.BaseModelTester):
     files_for_diff = [r"decoders/upernet/", r"base/"]
 
     default_batch_size = 2
+
+    @pytest.mark.torch_export
+    def test_torch_export(self):
+        super().test_torch_export(eps=1e-3)