diff --git a/baselines/common/atari_wrappers.py b/baselines/common/atari_wrappers.py
index 2c9b8c6..585081f 100644
--- a/baselines/common/atari_wrappers.py
+++ b/baselines/common/atari_wrappers.py
@@ -221,7 +221,7 @@ class LazyFrames(object):
         return len(self._force())
 
     def __getitem__(self, i):
-        return self._force()[i]
+        return self._force()[..., i]
 
 def make_atari(env_id, max_episode_steps=None):
     env = gym.make(env_id)