I’m using a Python Code for the same. A snippet of which I’m sharing below -
       for _ in range(epochs):
            self.model.train()
            train_loss = 0
            for idx, data in tqdm(enumerate(self.train_loader)):
                self.optimizer.zero_grad()
                output = self.model(input_ids = data["input_ids"], attention_mask = data["attention_mask"], lm_labels = data["lm_labels"])
                loss, prediction_scores = output[:2]  
                train_loss += loss.item()
                loss.backward()
                self.optimizer.step()
                if((idx % 1000) == 0):
                    print("loss: ", loss.item(), " train_loss: ", train_loss/(idx+1))

As the name of each variable suggests, I’ve PegasusForConditionalGeneration in the variable self.model, Adam Optimizer in self.optimizer and self.train_loader is of type DataLoader.