def forward(self, x):
features = self.features(x)
out = F.relu(features, inplace=True)
out = F.adaptive_avg_pool2d(out, (1, 1))
out = torch.flatten(out, 1)
label_out = torch.sigmoid(self.classifier_label(out))
rank_input = label_out if self.cat else out
rank_out = F.softmax(self.rank_classifier(rank_input), dim=1)
return label_out, rank_out
def pool_and_flatten(x):
x = F.adaptive_avg_pool2d(x, (1, 1))
return torch.flatten(x, 1)
