Project: 024_Progressive_distillation_induces_an_implicit_curri | Function: compute_loss (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes KL divergence loss between student and teacher outputs for progressive distillation. This implements the core distillation mechanism using knowledge transfer from teacher to student through KL divergence.

**KEY TERMS:** {}

**CONTEXT CODE:** import os 
from glob import glob 
import time
import pdb 

class KLTrainer_progressive(Trainer):
    
    def __init__(self, ref_model, temperature, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ref_model = ref_model.to(self.model.device)  #here we say that we have another model to refer to
        self.ref_config = self.ref_model.config
        self.ignore_index = -100
        self.temperature = temperature
        self.speed_factor = self.args.speed_factor
        self.refresh_frequency = self.args.refresh_frequency
        self.curr_alpha = 0.0
        self.ref_created_ckpts = []
      
        self.ref_model.requires_grad = False
        self.ref_model_ckpt_dir = self.args.ref_model_ckpt_dir

        self.distill_version = self.args.distill_version
        self.distill_topk = self.args.distill_topk

        self.step_cnt = 0
        self.logging_steps = self.args.logging_steps
        
        all_ckpts = glob(os.path.join(self.ref_model_ckpt_dir, 'checkpoint-*'))
        self.all_ckpts = {int(x.split('-')[-1]):x for x in all_ckpts}

        self.eval_acc = 0.
        self.eval_num = 0.
        self.eval_start = False
        
    def accuracy_eval(self, model, inputs, labels, padding_mask):
        if len(labels.shape) == 3:
            labels = labels.squeeze()
            padding_mask = padding_mask.squeeze()

        outputs = model(**inputs)
        #apply_softmax
        preds = outputs.logits.detach()
        accs = (preds[..., :-1, :].argmax(axis=-1) == labels).type(torch.float) * (1. - (padding_mask).type(torch.float))
        num_active_elements = padding_mask.numel() - padding_mask.long().sum()
        acc = accs.sum() / (1e-10 + num_active_elements) 

        self.eval_acc += acc.item() 
        self.eval_num += 1
        return acc.item()

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def training_step(self, model, inputs) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.

        Return:
            `torch.Tensor`: The tensor with training loss on this batch.
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    # custom log to include the step number
    def log(self, logs: Dict[str, float]) -> None:
        """
        Log `logs` on the various objects watching training.

        Subclass and override this method to inject custom behavior.

        Args:
            logs (`Dict[str, float]`):
                The values to log.
        """
        if self.state.epoch is not None:
            logs["epoch"] = round(self.state.epoch, 2)

        output = {**logs, **{"step": self.step_cnt}}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)
        


Complete the function implementation replacing "raise NotImplementedError".