Loss goes to NaN while using fp16 when working with arcface_torch and throws an error during validation

Hello
I am trying to train an arcface model on my custom dataset with 240226 ids. Here is my config file

```
from easydict import EasyDict as edict

config = edict()
config.margin_list = (1.0, 0.0, 0.4)
config.network = "r50"
config.resume = False
config.output = None
config.embedding_size = 512
config.sample_rate = 1.0
config.fp16 = True
config.momentum = 0.9
config.weight_decay = 1e-4
config.batch_size = 192
config.lr = 0.001
config.verbose = 2000
config.dali = False

config.rec = "/workspace/awi_facial_recognition/"
#config.use_pretrained = True
#config.pretrained_model_path = "/workspace/16backbone.pth"
config.num_classes = 240226
config.num_image = 12558871
config.num_epoch = 20
config.warmup_epoch = 0
config.val_targets = ['lfw', 'cfp_fp', 'agedb_30', 'calfw.bin']
```
After training for a couple of hours, I receive an error during validation. The error is regarding Nan values in the embedding. My loss also goes to nan before this error occurs. Can anyone please help on this?

Edit: There are several validation steps that take place before the error occurs. I sort of have an idea that there is something I need to do with either batch size or learning rate. However I would love to know some others ways I can solve it.
This error occurs before a whole epoch is completed. Please help me if possible.