//  [python]
def forward(self, features, gt_segm=None):
        batch_size = features.shape[0]

        init_pose = self.init_pose.expand(batch_size, -1)  # N, Jx6
        init_shape = self.init_shape.expand(batch_size, -1)
        init_cam = self.init_cam.expand(batch_size, -1)

        if self.use_position_encodings:
            features = torch.cat((features, self.pos_enc.repeat(features.shape[0], 1, 1, 1)), 1)

        output = {}

        ############## 2D PART BRANCH FEATURES ##############
        part_feats = self._get_2d_branch_feats(features)

        ############## GET PART ATTENTION MAP ##############
        part_attention = self._get_part_attention_map(part_feats, output)

        ############## 3D SMPL BRANCH FEATURES ##############
        smpl_feats = self._get_3d_smpl_feats(features, part_feats)

        ############## SAMPLE LOCAL FEATURES ##############
        if gt_segm is not None:
            # logger.debug(gt_segm.shape)
            # import IPython; IPython.embed(); exit()
            gt_segm = t.interpolate(gt_segm.unsqueeze(1).float(), scale_factor=(1/4, 1/4), mode='nearest').long().squeeze(1)
            part_attention = t.nn.functional.one_hot(gt_segm.to('cpu'), num_classes=self.num_joints + 1).permute(0,3,1,2).float()[:,1:,:,:]
            part_attention = part_attention.to('cuda')
            # part_attention = t.interpolate(part_attention, scale_factor=1/4, mode='bilinear', align_corners=True)
            # import IPython; IPython.embed(); exit()
        point_local_feat, cam_shape_feats = self._get_local_feats(smpl_feats, part_attention, output)

        ############## GET FINAL PREDICTIONS ##############
        pred_pose, pred_shape, pred_cam = self._get_final_preds(
            point_local_feat, cam_shape_feats, init_pose, init_shape, init_cam
        )

        if self.use_coattention:
            for c in range(self.num_coattention_iter):
                smpl_feats, part_feats = self.coattention(smpl_feats, part_feats)
                part_attention = self._get_part_attention_map(part_feats, output)
                point_local_feat, cam_shape_feats = self._get_local_feats(smpl_feats, part_attention, output)
                pred_pose, pred_shape, pred_cam = self._get_final_preds(
                    point_local_feat, cam_shape_feats, pred_pose, pred_shape, pred_cam
                )

        if self.num_branch_iteration > 0:
            for nbi in range(self.num_branch_iteration):
                if self.use_soft_attention:
                    smpl_feats = self.branch_iter_3d_nonlocal(smpl_feats)
                    part_feats = self.branch_iter_2d_nonlocal(part_feats)
                else:
                    smpl_feats = self.branch_iter_3d_nonlocal(smpl_feats)
                    part_feats = smpl_feats

                part_attention = self._get_part_attention_map(part_feats, output)
                point_local_feat, cam_shape_feats = self._get_local_feats(smpl_feats, part_attention, output)
                pred_pose, pred_shape, pred_cam = self._get_final_preds(
                    point_local_feat, cam_shape_feats, pred_pose, pred_shape, pred_cam,
                )

        pred_rotmat = rot6d_to_rotmat(pred_pose).reshape(batch_size, 24, 3, 3)

        output.update({
            'pred_pose': pred_rotmat,
            'pred_cam': pred_cam,
            'pred_shape': pred_shape,
        })
        return output