diff --git a/tests/test_vec_check_nan.py b/tests/test_vec_check_nan.py
index 85dd884..a04410a 100644
--- a/tests/test_vec_check_nan.py
+++ b/tests/test_vec_check_nan.py
@@ -2,7 +2,7 @@ import gym
 from gym import spaces
 import numpy as np
 
-from stable_baselines.common.vec_env import DummyVecEnv, VecCheckNan
+from stable_baselines3.common.vec_env import DummyVecEnv, VecCheckNan
 
 
 class NanAndInfEnv(gym.Env):