//  [python]
def _decode(self, sample, styled_hidden_states): 
        transformer_out = super()._decode(sample, styled_hidden_states)
        B, T = transformer_out.shape[:2]

        vector_idx = 0
        if self.predict_jaw:
            rotation_rep_size_ = rotation_rep_size(self.rotation_representation)
            jaw_pose = transformer_out[..., :rotation_rep_size_].view(B*T, -1)
            # ensure orthonogonality of the 6D rotation
            if self._rotation_representation() in ["6d", "6dof"]:
                jaw_pose = matrix_to_rotation_6d(rotation_6d_to_matrix(jaw_pose))
            elif self._rotation_representation() in ['mat', 'matrix']:
                jaw_pose = jaw_pose.view(-1, 3, 3) 
                U, S, Vh = torch.linalg.svd(jaw_pose) # do SVD to ensure orthogonality
                jaw_pose = U.contiguous()
            vector_idx += rotation_rep_size_
            
        else: 
            jaw = sample["gt_jaw"][:, :T, ...]
            jaw_pose = jaw.view(B*T, -1)
        if self.predict_exp: 
            expression_params = transformer_out[..., vector_idx:].view(B*T, -1)
            vector_idx += self.flame_config.n_exp
        else: 
            exp = sample["gt_exp"][:, :T, ...]
            expression_params = exp.view(B*T, -1)
        
        assert vector_idx == transformer_out.shape[-1]

        sample["predicted_exp"] = expression_params.view(B,T, -1)
        sample["predicted_jaw"] = jaw_pose.view(B,T, -1)
        if "gt_shape" not in sample.keys():
            template = sample["template"]
            assert B == 1, "Batch size must be 1 if we want to change the template inside FLAME"
            self.flame.v_template = template.squeeze(1).view(-1, 3)
            shape_params = torch.zeros((B,T, self.flame_config.n_shape), device=transformer_out.device)
        else: 
            # sample["gt_shape"].shape = (B, n_shape) -> add T dimension and repeat
            shape_params = sample["gt_shape"][:, None, ...].repeat(1, T, 1)

        flame_jaw_pose = convert_rot(jaw_pose, self.rotation_representation, self.flame_rotation_rep)
        pose_params = self.flame.eye_pose.expand(B*T, -1)
        pose_params = torch.cat([ pose_params[..., :3], flame_jaw_pose], dim=-1)

        # with torch.no_grad():
        #     # vertice_neutral, _, _ = self.flame.forward(shape_params[0:1, ...], torch.zeros_like(expression_params[0:1, ...])) # compute neutral shape
        #                 # vertice_neutral, _, _ = self.flame.forward(shape_params[:, 0, ...], torch.zeros_like(expression_params[0:1, ...])) # compute neutral shape
        #     zero_exp = torch.zeros((B, expression_params.shape[1]), dtype=shape_params.dtype, device=shape_params.device)
        #     # compute neutral shape for each batch (but not each frame, unnecessary)
        #     vertices_neutral, _, _ = self.flame.forward(shape_params[:, 0, ...], zero_exp) # compute neutral shape
        #     vertices_neutral = vertices_neutral.contiguous().view(vertices_neutral.shape[0], -1)[:, None, ...]

        vertices_neutral = self._neutral_shape(B, expression_params.shape[1], shape_params)
        # vertices_neutral = vertices_neutral.expand(B, T, -1, -1)

        shape_params = shape_params.view(B * T, -1)
        vertices_out, _, _ = self.flame(shape_params, expression_params, pose_params)
        vertices_out = vertices_out.contiguous().view(B, T, -1)
        vertex_offsets = vertices_out - vertices_neutral # compute the offset that is then added to the template shape
        vertex_offsets = vertex_offsets.view(B, T, -1)
        return vertex_offsets