@Suteja Kanuri​ Update - I was able to get it to work by upgrading to a g4dn.12xlarge node (4 gpus).

However, the code in 02-Data-preparation to apply sshleifer/distilbart-cnn-12-6 model for a summarization task failed with the more powerful node (while it worked fine with just a single GPU). Do you have any suggestions there?

I set repartition to 4 since there were 4 GPUs. docs_limit_df has 4 rows.

torch.cuda.empty_cache()
 
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device_map="auto")
 
docs_limit_df = docs_limit_df.repartition(4).withColumn("text_short", summarize_all("text"))

The error I got was

"PythonException: 'RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0!', from <command-434574176370212>, line 8. Full traceback below:"