@karthikgali Something is confusing about this post for me.  You wrote

We have finetuned distill-pegasus-cnn-16-4 2

But I can see in your code that you are loading
model_name = “sshleifer/distill-pegasus-cnn-16-4”
If you fine-tuned it, wouldn’t you load your fine-tuned model, and not the model that has already been tuned by sshleifer?
Put another way, see here: https://huggingface.co/transformers/model_doc/pegasus.html#checkpoints

All the checkpoints are fine-tuned for summarization, besides pegasus-large, whence the other checkpoints are fine-tuned

My understanding from this is that, if we, as the newbie user, have some data we want to use with Pegasus, we should do this:

Start with pegasus-large: https://huggingface.co/google/pegasus-large 2

Fine tune it on our own data
Use the pytorch_model.bin output from this fine tuning process to run inference on our own data.

Note – I am not suggesting you are doing anything wrong. I’m just working with pegasus also, and reading your message makes me think that I’m maybe getting something wrong.
Also – I am finding, when fine-tuning Pegasus, using pegasus-large, that the RAM requirements for even just a batch size of 1 are so extreme, that a Nvidia card with 16GB of memory is required… just to run the batch size of 1! So at this point I am thinking that maybe my training will run better on the CPU, using a machine with a huge amount of ram… like 512GB of ram… as this seems to allow a much bigger batch size, like up to 64 or 128.