examples/research_projects/visual_bert/modeling_frcnn.py/ROIOutputs/_predict_attrs
class ROIOutputs:    def _predict_attrs(self, attr_logits, preds_per_image):
        attr_logits = attr_logits[..., :-1].softmax(-1)
        attr_probs, attrs = attr_logits.max(-1)
        return attr_probs.split(preds_per_image, dim=0), attrs.split(preds_per_image, dim=0)