| 17 | |
| 18 | |
| 19 | class Wrapper(nn.Module): |
| 20 | def __init__(self, model): |
| 21 | super(Wrapper, self).__init__() |
| 22 | self.model = model |
| 23 | self.softmax = nn.Softmax(dim=1) |
| 24 | |
| 25 | def forward(self, x): |
| 26 | x = self.model(x) |
| 27 | x = self.softmax(x) |
| 28 | return x |
| 29 | |
| 30 | |
| 31 | class KestrelSolver(ClsSolver): |