diff --git a/5.4-visualizing-what-convnets-learn.ipynb b/5.4-visualizing-what-convnets-learn.ipynb
index 8697872303..bbeb043459 100644
--- a/5.4-visualizing-what-convnets-learn.ipynb
+++ b/5.4-visualizing-what-convnets-learn.ipynb
@@ -1034,7 +1034,7 @@
     "# We multiply each channel in the feature map array\n",
     "# by \"how important this channel is\" with regard to the elephant class\n",
     "for i in range(512):\n",
-    "    conv_layer_output_value[:, :, i] *= pooled_grads_value[i]\n",
+    "    conv_layer_output_value[:, :, :, i] *= pooled_grads_value[i]\n",
     "\n",
     "# The channel-wise mean of the resulting feature map\n",
     "# is our heatmap of class activation\n",
@@ -1149,4 +1149,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file