diff --git a/columbus/env.py b/columbus/env.py
index 2b6d02d..dc4b17d 100644
--- a/columbus/env.py
+++ b/columbus/env.py
@@ -323,7 +323,7 @@ class ColumbusEnv(gym.Env):
             c = 0
             for i in range(resolution):
                 for j in range(resolution):
-                    v = V[c]
+                    v = V[c].item()
                     c += 1
                     col = [int((1-v)*color_depth),
                            int(v*color_depth), 0, color_depth]