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
        """
        option_num = self.task.fewshot_task.class_num
        assert sample["targets"].size(0) % option_num == 0
        sample_size = sample["targets"].size(0) // option_num
        loss_mask = sample["net_input"]["gpt_loss_mask"][:, 1:]
        context_len = sample["context_len"]

        if len(loss_mask) == loss_mask.int().sum():
            # single-token classification tasks
            shape = sample["net_input"]["src_tokens"].shape
            src_tokens = sample["net_input"]["src_tokens"].view(len(loss_mask)//option_num, option_num, -1)[:, 0]
            net_output, extra = model(
                src_tokens
            )
            net_output = net_output.unsqueeze(1).expand(-1, option_num, -1, -1).reshape(len(loss_mask), len(src_tokens[0]), -1)
        else: 
            # multi-token classification tasks and multi-choice tasks
            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(-1)
        loss = loss / loss_mask.int().sum(-1)  # (bsz, option_num)

        fewshot_labels = sample["targets"].view(-1)
        pred_label = torch.argmax(loss.view(-1, option_num), dim=1)
        target_label = fewshot_labels.view(-1, option_num)[:, 0]

        loss = loss.view(-1, option_num)
        tmp_loss = 0
        for i in range(sample_size):
            tmp_loss += loss[i][target_label[i]]
        loss = tmp_loss

        logging_output = {}

        logging_output.update(
            {
                "loss": -loss.sum().data,
                "ntokens": sample["ntokens"],
                "nsentences": sample_size,
                "sample_size": sample_size,
                "ncorrect": (pred_label == target_label).sum(),
                "npos": (target_label == 0).sum(),
                "nneg": (target_label == 1).sum(),
            }
        )
        
        # ! ICL analysis
        def process_hiddens_classification(hiddens):
            # single-token classification tasks
            hiddens = torch.stack(hiddens, dim=0)  # (n_layers, len, bsz, hidden)
            hiddens = hiddens.transpose(1, 2)  # (n_layers, bsz, len, hidden)
            record_hiddens = hiddens[:, 0, -1, :]  # (n_layers, hidden)
            return record_hiddens


        def process_attention_map_classification(attn_map, c_len):
            # single-token classification tasks
            attn_map = torch.stack(attn_map, dim=0)  # (n_layers, n_heads, len, len)
            record_attn_map = attn_map[:, :, -1, c_len:]  # (n_layers, n_heads, query_len)
            record_attn_map_ctx = attn_map[:, :, -1, :c_len]  # (n_layers, n_heads, context_len)
            # record_attn_map = record_attn_map.sum(dim=-2)  # (n_layers, query_len)
            record_attn_map_ctx = record_attn_map_ctx.sum(dim=-2)  # (n_layers, context_len)
            return record_attn_map, record_attn_map_ctx


        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_self_attn_out_hiddens = None
        record_trm_out_hiddens = None
        record_attn_map = None
        record_attn_map_ctx = None
        attn_q = None
        
        record_info = None
        
        if len(loss_mask) == loss_mask.int().sum():
            # single-token classification tasks
            record_self_attn_out_hiddens = process_hiddens_classification(extra['self_attn_out_hiddens'])
            qkv_val = extra['qkv_val']
            attn_map = [item['record_attn_map'] for item in qkv_val]
            record_attn_map, record_attn_map_ctx = process_attention_map_classification(attn_map, context_len[0])
            attn_q = [item['q'] for item in qkv_val]
            record_attn_q = process_attention_q_classification(attn_q)
        
        record_info = {
            'self_attn_out_hiddens': record_self_attn_out_hiddens.tolist(),
            'attn_map': record_attn_map.tolist(),
            'attn_map_ctx': record_attn_map_ctx.tolist(),
            'attn_q': record_attn_q,
            'gold_label': target_label[0].item(),
            'pred_label': pred_label[0].item(),
        }
        
        return loss, sample_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_eval.py