Fixed LeNet architecture not working (#23)

* Fixed LeNet architecture not working * Apply suggestions from code review: Roll back name for LeNet Co-authored-by: btwardow <[email protected]>
mmasana · Sep 15, 2022 · e2176a3 · e2176a3
1 parent f653d6c
commit e2176a3
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/networks/lenet.py b/src/networks/lenet.py
@@ -2,7 +2,7 @@
 import torch.nn.functional as F
 
 
-class LeNet(nn.Module):
+class LeNetArch(nn.Module):
     """LeNet-like network for tests with MNIST (28x28)."""
 
     def __init__(self, in_channels=1, num_classes=10, **kwargs):
@@ -28,3 +28,9 @@ def forward(self, x):
         out = F.relu(self.fc2(out))
         out = self.fc(out)
         return out
+
+
+def LeNet(pretrained=False, **kwargs):
+    if pretrained:
+        raise NotImplementedError
+    return LeNetArch(**kwargs)