diff --git a/tests/test_sb3_onnx_export.py b/tests/test_sb3_onnx_export.py
index baeea9b9..362a77ed 100644
--- 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"