Hi Richard sorry I didn't see your message, I assume you figured it out but for future readers here is the full script (with loader definition):

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