def inner_forward(self, start, end):
        '''
        Re-compute the forward pass from the checkpointed model parameters.
        '''
        if start % self.ws == self.r:
            self.thetas.append(self.model.get_params_vec())
        
        for i in range(start + 1, end):
            # the i-th update, use data with index i-1
            # i is the index of the model after the i-th update
            torch.cuda.synchronize()
            st = time()
            batch = self.checkpointing.load_h5_data(i-1)
            
            logging_loss = 0
            for ii in range(self.gacc):
                mini_batch = {k: v[ii*self.gl_bs:(ii+1)*self.gl_bs][self.r*self.bs:(self.r+1)*self.bs] for k, v in batch.items()}
                loss, _ = self.model.compute_loss(**mini_batch)
                loss = loss / (self.ws * self.gacc)
                loss.backward()
                dist.all_reduce(loss, op=dist.ReduceOp.SUM)
                logging_loss += loss.item()

            for param in self.model.parameters():
                dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)

            self.optimizer.step()
            self.lr_scheduler.step()
            self.optimizer.zero_grad()

            if i % self.ws == self.r:
                self.thetas.append(self.model.get_params_vec())
            
            torch.cuda.synchronize()
            inner_fw_time = time() - st
            
            if i % self.args.log_interval == 0:
                log_str = "IF: Af. {}-th Forw. Upd., Loss: {:.4f}, IF_time: {:.2f}".format(i, logging_loss, inner_fw_time)
                print_rank(log_str)

        Re-compute the forward pass from the checkpointed model parameters.
        
data_selection/pmp_solver/trainer.py