Hey Harini sorry I missed your message, I assume you figured it out already but for future readers the full code can be found here

https://github.com/jamescalam/pytorch-mps/blob/main/code/02mps_training_bert_head.ipynb