diff --git a/setup.py b/setup.py
index b371643..2949718 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@ setup(name='torchy_baselines',
       install_requires=[
           'gym[classic_control]>=0.10.9',
           'numpy',
-          'torch>=1.2.0' # torch>=1.2.0+cpu
+          'torch>=1.2.0'
       ],
       extras_require={
         'tests': [
diff --git a/tests/test_td3.py b/tests/test_td3.py
index ad2a116..a6fdb15 100644
--- a/tests/test_td3.py
+++ b/tests/test_td3.py
@@ -5,8 +5,7 @@ import gym
 from torchy_baselines import TD3
 
 def test_pendulum():
-    env = gym.make("Pendulum-v0")
-    model = TD3('MlpPolicy', env, policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1)
+    model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1)
     model.learn(total_timesteps=500, eval_freq=100)
     model.save("test_save")
     model.load("test_save")
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 4916f0a..2f5bc78 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -34,6 +34,9 @@ class BaseRLModel(object):
         self.params = None
 
         if env is not None:
+            if env is not None:
+                if isinstance(env, str):
+                    env = gym.make(env)
             self.env = env
             self.n_envs = 1
             self.observation_space = env.observation_space