| 123 | |
| 124 | # Multiclass problem |
| 125 | class NeuralNet2(nn.Module): |
| 126 | def __init__(self, input_size, hidden_size, num_classes): |
| 127 | super(NeuralNet2, self).__init__() |
| 128 | self.linear1 = nn.Linear(input_size, hidden_size) |
| 129 | self.relu = nn.ReLU() |
| 130 | self.linear2 = nn.Linear(hidden_size, num_classes) |
| 131 | |
| 132 | def forward(self, x): |
| 133 | out = self.linear1(x) |
| 134 | out = self.relu(out) |
| 135 | out = self.linear2(out) |
| 136 | # no softmax at the end |
| 137 | return out |
| 138 | |
| 139 | model = NeuralNet2(input_size=28*28, hidden_size=5, num_classes=3) |
| 140 | criterion = nn.CrossEntropyLoss() # (applies Softmax) |
no outgoing calls
no test coverage detected