| 2 | import torch.nn as nn |
| 3 | |
| 4 | class MultiClassNet(nn.Module): |
| 5 | def __init__(self, NUM_FEATURES, NUM_CLASSES, HIDDEN_FEATURES): |
| 6 | super().__init__() |
| 7 | self.lin1 = nn.Linear(NUM_FEATURES, HIDDEN_FEATURES) |
| 8 | self.lin2 = nn.Linear(HIDDEN_FEATURES, NUM_CLASSES) |
| 9 | self.log_softmax = nn.LogSoftmax(dim=1) |
| 10 | |
| 11 | def forward(self, x): |
| 12 | x = self.lin1(x) |
| 13 | x = torch.sigmoid(x) |
| 14 | x = self.lin2(x) |
| 15 | x = self.log_softmax(x) |
| 16 | return x |
no outgoing calls
no test coverage detected