Update test_sb3_onnx_export.py

Updates onnx export method name
edbeeching · Aug 22, 2024 · aa4e216 · aa4e216
1 parent 95cb1ad
commit aa4e216
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_sb3_onnx_export.py b/tests/test_sb3_onnx_export.py
@@ -24,7 +24,7 @@
 def test_pytorch_vs_onnx(env_name, port):
     from stable_baselines3 import PPO
 
-    from godot_rl.wrappers.onnx.stable_baselines_export import export_ppo_model_as_onnx, verify_onnx_export
+    from godot_rl.wrappers.onnx.stable_baselines_export import export_model_as_onnx, verify_onnx_export
     from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv
 
     env_path = f"examples/godot_rl_{env_name}/bin/{env_name}.x86_64"