fix unittests

microsoft · Nov 20, 2023 · e3e1001 · e3e1001
1 parent 7537ca7
commit e3e1001
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py
@@ -18,7 +18,10 @@
 
 
 def get_opsets_model(filename):
-    onx = onnx.load(filename)
+    if isinstance(filename, onnx.ModelProto):
+        onx = filename
+    else:
+        onx = onnx.load(filename)
     return {d.domain: d.version for d in onx.opset_import}
 
 
@@ -1005,9 +1008,9 @@ def test_save_ort_format():
         assert os.path.exists(os.path.join(temp_dir, "optimizer_model.onnx"))
         assert os.path.exists(os.path.join(temp_dir, "optimizer_model.ort"))
         base_opsets = get_opsets_model(base_model)
-        training_opsets = get_opsets_model("training_model.onnx")
-        eval_opsets = get_opsets_model("eval_model.onnx")
-        optimizer_opsets = get_opsets_model("optimizer_model.onnx")
+        training_opsets = get_opsets_model(os.path.join(temp_dir, "training_model.onnx"))
+        eval_opsets = get_opsets_model(os.path.join(temp_dir, "eval_model.onnx"))
+        optimizer_opsets = get_opsets_model(os.path.join(temp_dir, "optimizer_model.onnx"))
         if base_opsets[""] != training_opsets[""]:
             raise AssertionError(f"Opsets mismatch {base_opsets['']} != {training_opsets['']}.")
         if base_opsets[""] != eval_opsets[""]: