Fix kernel dimensions for LeNet model code example (#2192)

kiersten-stokes · subramen · Svetlana Karslioglu · web-flow · commit 3764368aa14a · 2023-02-14T14:57:27.000-08:00
* Change kernel to 5x5 in  1st Conv2d layer in model init

Signed-off-by: Kiersten Stokes &lt;kierstenstokes@gmail.com&gt;

* Change kernel to 5x5 in 2nd Conv2d layer in model init

* Fix dimensions of 1st Linear layer to match new expected size
---------
Signed-off-by: Kiersten Stokes &lt;kierstenstokes@gmail.com&gt;
Co-authored-by: Suraj Subramanian &lt;5676233+suraj813@users.noreply.github.com&gt;
Co-authored-by: Svetlana Karslioglu &lt;svekars@fb.com&gt;
diff --git a/beginner_source/introyt/introyt1_tutorial.py b/beginner_source/introyt/introyt1_tutorial.py
@@ -176,12 +176,12 @@ class LeNet(nn.Module):
 
     def __init__(self):
         super(LeNet, self).__init__()
-        # 1 input image channel (black & white), 6 output channels, 3x3 square convolution
+        # 1 input image channel (black & white), 6 output channels, 5x5 square convolution
         # kernel
-        self.conv1 = nn.Conv2d(1, 6, 3)
-        self.conv2 = nn.Conv2d(6, 16, 3)
+        self.conv1 = nn.Conv2d(1, 6, 5)
+        self.conv2 = nn.Conv2d(6, 16, 5)
         # an affine operation: y = Wx + b
-        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension
         self.fc2 = nn.Linear(120, 84)
         self.fc3 = nn.Linear(84, 10)
 
