fix(LeNet5): 批量大小为1时LeNet5前向操作中C5->F6步骤出错

deep-learning-algorithm · May 28, 2019 · 48b5a46 · 48b5a46
1 parent 844d749
commit 48b5a46
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nn/nets.py b/nn/nets.py
@@ -109,7 +109,7 @@ def forward(self, inputs):
         x = self.maxPool2(x)
         x = self.relu3(self.conv3(x))
         # (N, C, 1, 1) -> (N, C)
-        x = x.squeeze()
+        x = x.reshape(x.shape[0], -1)
         x = self.relu4(self.fc1(x))
         x = self.fc2(x)