From e2176a3eb2ade9fbb4b284614afcb415ff9f41fa Mon Sep 17 00:00:00 2001 From: Marc Masana Date: Thu, 15 Sep 2022 10:57:22 +0200 Subject: [PATCH] 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 --- src/networks/lenet.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/networks/lenet.py b/src/networks/lenet.py index 2b4f012..b556cb3 100644 --- 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)