def forward(self, model, sample, reduce=True, split="valid"):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        loss_mask = sample["net_input"]["gpt_loss_mask"][:, 1:]
        sample_size = sample["net_input"]["src_tokens"].size(0)
        net_output, extra = model(
            sample["net_input"]["src_tokens"]
        )
        net_output = net_output[:, :-1, :]
        net_output = (net_output, extra)
        targets = sample["net_input"]["src_tokens"][:, 1:].unsqueeze(-1)

        lprobs = model.get_normalized_probs(net_output, log_probs=True)
        loss = torch.gather(lprobs, -1, targets).squeeze(-1) * (loss_mask != False).int()
        loss = -loss.sum()

        optim_size = loss_mask.int().sum()

        logging_output = {}

        logging_output.update(
            {
                "loss": loss.data,
                "ntokens": sample["ntokens"],
                "nsentences": sample_size,
                "sample_size": optim_size,
            }
        )


        # ! ICL analysis
        def process_attention_q_classification(attn_q):
            # single-token classification tasks
            attn_q = torch.stack(attn_q, dim=0)  # (n_layers, n_token, bzs (1), q_hidden_dim)
            record_attn_q = attn_q[:, :, -1, :]  # (n_layers, n_token, q_hidden_dim)
            return record_attn_q

        record_attn_q = None

        record_info = None

        qkv_val = extra['qkv_val']
        attn_q = [item['q'] for item in qkv_val]
        record_attn_q = process_attention_q_classification(attn_q)

        record_info = {
            'attn_q': record_attn_q,
        }

        return loss, optim_size, logging_output, record_info
Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        
understand_icl/struprompting/criterions/fs_ft.py