diff --git a/style-transfer/Style_Transfer_Solution.ipynb b/style-transfer/Style_Transfer_Solution.ipynb
index c2e15802fe..9de1c08632 100644
--- a/style-transfer/Style_Transfer_Solution.ipynb
+++ b/style-transfer/Style_Transfer_Solution.ipynb
@@ -358,10 +358,10 @@
     "    \"\"\"\n",
     "    \n",
     "    # get the batch_size, depth, height, and width of the Tensor\n",
-    "    _, d, h, w = tensor.size()\n",
+    "    b, d, h, w = tensor.size()\n",
     "    \n",
     "    # reshape so we're multiplying the features for each channel\n",
-    "    tensor = tensor.view(d, h * w)\n",
+    "    tensor = tensor.view(b * d, h * w)\n",
     "    \n",
     "    # calculate the gram matrix\n",
     "    gram = torch.mm(tensor, tensor.t())\n",
@@ -651,7 +651,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python [default]",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -665,7 +665,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.4"
+   "version": "3.7.3"
   }
  },
  "nbformat": 4,