fix

onnx · Feb 27, 2025 · df637d4 · df637d4
1 parent e37520a
commit df637d4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_sklearn_nearest_neighbour_converter.py b/tests/test_sklearn_nearest_neighbour_converter.py
@@ -1760,10 +1760,10 @@ def test_knn_imputer_one_nan(self):
         input_data = data.astype(np.float32)
         initial_type = [("float_input", FloatTensorType([None, data.shape[1]]))]
         onnx_model = convert_sklearn(imputer, initial_types=initial_type)
-        with open("test_knn_imputer_one_nan.onnx", "wb") as f:
-            f.write(onnx_model.SerializeToString())
+        # with open("test_knn_imputer_one_nan.onnx", "wb") as f:
+        #     f.write(onnx_model.SerializeToString())
 
-        sess = rt.InferenceSession("test_knn_imputer_one_nan.onnx")
+        sess = rt.InferenceSession(onnx_model.SerializeToString())
         input_name = sess.get_inputs()[0].name
         output_name = sess.get_outputs()[0].name
         res = sess.run([output_name], {input_name: input_data})