use pretrained weights for torch.compile

gau-nernst · Jul 7, 2023 · a10fe37 · a10fe37
1 parent b64adb4
commit a10fe37
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_backbones.py b/tests/test_backbones.py
@@ -68,6 +68,6 @@ def test_jit_trace(self, name: str, inputs: Tensor):
 
     @pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile() is not available")
     def test_compile(self, name: str, inputs: Tensor):
-        m = getattr(backbones, name)().eval()
+        m = getattr(backbones, name)(pretrained=True).eval()
         m_compiled = torch.compile(m)
         torch.testing.assert_close(m(inputs), m_compiled(inputs))