diff --git a/neural_style.lua b/neural_style.lua
index 0bb4f08..a2ab239 100644
--- a/neural_style.lua
+++ b/neural_style.lua
@@ -149,9 +149,8 @@ local function main(params)
     end
     net:add(tv_mod)
   end
-  for i = 1, #cnn do
+  for i,layer in ipairs(cnn.modules) do
     if next_content_idx <= #content_layers or next_style_idx <= #style_layers then
-      local layer = cnn:get(i)
       local name = layer.name
       local layer_type = torch.type(layer)
       local is_pooling = (layer_type == 'cudnn.SpatialMaxPooling' or layer_type == 'nn.SpatialMaxPooling')