(self, weight=None, size_average=True, ignore_index=255)
| 40 | |
| 41 | class CrossEntropyLoss2d(nn.Module): |
| 42 | def __init__(self, weight=None, size_average=True, ignore_index=255): |
| 43 | super(CrossEntropyLoss2d, self).__init__() |
| 44 | self.nll_loss = nn.NLLLoss2d(weight, size_average, ignore_index) |
| 45 | |
| 46 | def forward(self, inputs, targets): |
| 47 | return self.nll_loss(F.log_softmax(inputs), targets) |