def compute_proxy_ct(self, lam_param, params, buffers):
        '''
        Compute the contribution of the proxy data to the gradient of the loss w.r.t. the model parameters.
        '''
        proxy_sampler = DistributedSampler(
            self.proxy_dataset, shuffle=False, drop_last=False, rank=self.r, num_replicas=self.ws)
        proxy_dataloader = DataLoader(
            self.proxy_dataset, sampler=proxy_sampler, batch_size=self.px_bs,
            num_workers=self.args.num_workers, collate_fn=self.train_dataset.collate, drop_last=False)
        
        all_ct = []
        for i, (model_batch, no_model_batch) in enumerate(proxy_dataloader):
            if i % 10 == 0:
                print_rank("Computing contribution on proxy data: {}/{}".format(i, len(proxy_dataloader)))
            self.proxy_dataset.move_to_device(model_batch, no_model_batch, self.device)
            batch = {**model_batch, **no_model_batch}
            ct = jvp_batch(self.model, batch, lam_param, params, buffers)
            ct = all_gather(ct, dim=1, op="stack").view(-1)
            ct = self.args.lr * ct
            all_ct.append(ct)
        all_ct = torch.cat(all_ct, dim=0)
        all_ct = all_ct[:len(self.proxy_dataset)]
        
        return all_ct

        Compute the contribution of the proxy data to the gradient of the loss w.r.t. the model parameters.
        
data_selection/pmp_solver/trainer.py