Spaces:
Sleeping
Sleeping
rzimmerdev
commited on
Commit
•
7dc7452
1
Parent(s):
2262103
fix: Removed last activation layer from LeNet model
Browse files- src/models.py +2 -1
src/models.py
CHANGED
@@ -30,7 +30,8 @@ class CNN(nn.Module):
|
|
30 |
classifier_layers.append(
|
31 |
nn.Linear(in_features=layer[0], out_features=layer[1])
|
32 |
)
|
33 |
-
|
|
|
34 |
|
35 |
self.feature_extractor = nn.Sequential(*feature_layers)
|
36 |
self.classifier = nn.Sequential(*classifier_layers)
|
|
|
30 |
classifier_layers.append(
|
31 |
nn.Linear(in_features=layer[0], out_features=layer[1])
|
32 |
)
|
33 |
+
if idx < len(self.classifier_activations) - 1:
|
34 |
+
classifier_layers.append(self.classifier_activations[idx]())
|
35 |
|
36 |
self.feature_extractor = nn.Sequential(*feature_layers)
|
37 |
self.classifier = nn.Sequential(*classifier_layers)
|