def classifier(channel, pool, out):
    layers = [nn.Conv2d(channel, out, kernel_size=1), nn.AvgPool2d(pool), ViewFlatten()]
    return nn.Sequential(*layers)

out = classifier(input)
1) top_prob, top_label = torch.topk(F.softmax(out, dim=1), 1)
2) top_prob, top_label = out.max(1)

which one is right to get top probability and top label (predicted label) ?