Added hidden size as parameter
@@ -28,15 +28,15 @@ import torch.nn as nn
class SimpleMLP(nn.Module):
- def __init__(self, num_classes=10, input_size=28*28):
+ def __init__(self, num_classes=10, input_size=28*28, hidden_size=512):
super(SimpleMLP, self).__init__()
self.features = nn.Sequential(
- nn.Linear(input_size, 512),
+ nn.Linear(input_size, hidden_size),
nn.ReLU(inplace=True),
nn.Dropout(),
)
- self.classifier = nn.Linear(512, num_classes)
+ self.classifier = nn.Linear(hidden_size, num_classes)
self._input_size = input_size
def forward(self, x):
