fix(test): dump and load onnx models

ise-uiuc · Mar 19, 2024 · 6e0f0e0 · 6e0f0e0
1 parent 971e397
commit 6e0f0e0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/torch/test_dump_load.py b/tests/torch/test_dump_load.py
@@ -45,7 +45,7 @@ def test_onnx_load_dump(tmp_path):
     # check oracle
     compare_two_oracle(oracle, loaded_testcase.oracle)
 
-    loaded_model = loaded_testcase.model.torch_model
+    loaded_model = loaded_testcase.model
     loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()}
     rerun_oracle = loaded_model.make_oracle()
     compare_two_oracle(oracle, rerun_oracle)
@@ -77,7 +77,7 @@ def test_bug_report_load_dump(tmp_path):
     # check oracle
     compare_two_oracle(oracle, loaded_testcase.oracle)
 
-    loaded_model = loaded_testcase.model.torch_model
+    loaded_model = loaded_testcase.model
     loaded_model.sat_inputs = {k: torch.from_numpy(v) for k, v in oracle.input.items()}
     rerun_oracle = loaded_model.make_oracle()
     compare_two_oracle(oracle, rerun_oracle)