examples/research_projects/visual_bert/modeling_frcnn.py/FastRCNNOutputLayers/forward
class FastRCNNOutputLayers:    def forward(self, roi_features):
        if roi_features.dim() > 2:
            roi_features = torch.flatten(roi_features, start_dim=1)
        scores = self.cls_score(roi_features)
        proposal_deltas = self.bbox_pred(roi_features)
        if self.use_attr:
            _, max_class = scores.max(-1)  # [b, c] --> [b]
            cls_emb = self.cls_embedding(max_class)  # [b] --> [b, 256]
            roi_features = torch.cat([roi_features, cls_emb], -1)  # [b, 2048] + [b, 256] --> [b, 2304]
            roi_features = self.fc_attr(roi_features)
            roi_features = nn.functional.relu(roi_features)
            attr_scores = self.attr_score(roi_features)
            return scores, attr_scores, proposal_deltas
        else:
            return scores, proposal_deltas