SyncBatchNorm causing NaN predictions during pertaining

[Shra1-25]

Hello everyone,

I am running the pretraining of SLIP model on ISIC pathology dataset and noticed some nan predictions of some images. The images have been appropriately normalised so there is no issue with the dataset.
If I remove the SyncBatchNorm layers, there are no issues of nan predictions but the model gradients are very small causing no leaning during the pretraining.
It is weird since the model trains well for some iterations and eventually fails after the epoch is 40% complete.

Any help/suggestions would be appreciated!