nakranivaibhav

@amyeroberts is this up for taking?

Alright @muellerzr i will look into it.

Okay so when i try to replicate this on a smaller model, bert-base-uncased ,there is no error. Does the model have to be big enough so that it gets sharded...