Spaces:
Sleeping
Sleeping
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)
|