Nan in output of model through several epoch of training

I am trying to train a face recognition model. Model consists of CNN and head - ArcFace (https://github.com/ronghuaiyang/arcface-pytorch)
I am using optimization levels opt_level = O3

My pipeline almost completely copies the code from the samples (https://github.com/NVIDIA/apex/tree/master/examples/imagenet)

global best_prec1, args
    args = parse()
    print("opt_level = {}".format(args.opt_level))
    print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
    print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
    print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
    RGB_MEAN = [0.5, 0.5, 0.5]
    RGB_STD = [0.5, 0.5, 0.5]
    cudnn.benchmark = True
    best_prec1 = 0
    if args.deterministic:
        cudnn.benchmark = False
        cudnn.deterministic = True
        torch.manual_seed(args.local_rank)
        torch.set_printoptions(precision=10)
    args.distributed = True
    if 'WORLD_SIZE' in os.environ:
        args.distributed = int(os.environ['WORLD_SIZE']) > 1
    args.gpu = 0
    args.world_size = 1
    if args.distributed:
        args.gpu = args.local_rank
        torch.cuda.set_device(args.gpu)
        torch.distributed.init_process_group(backend='nccl',
                                             init_method='env://')
        args.world_size = torch.distributed.get_world_size()
    assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."
    if args.channels_last:
        memory_format = torch.channels_last
    else:
        memory_format = torch.contiguous_format
    # create model
    model = IR_152([112,112])
    print('IR_152 generated')
    head = ArcFace(in_features = 256, out_features = 1607693, device_id = None) #
    print('HEAD generated')
    # sync bn
    import apex
    print("using apex synced BN")
    model = apex.parallel.convert_syncbn_model(model)
    head = apex.parallel.convert_syncbn_model(head)
    model = model.cuda().to(memory_format=memory_format)
    head = head.cuda().to(memory_format=memory_format)
    # Scale learning rate based on global batch size
    args.lr = args.lr*float(args.batch_size*args.world_size)/256.
    # args.lr = 0.1
    print('LR - ', args.lr)
    optimizer = torch.optim.SGD([{'params': model.parameters()}, {'params': head.parameters()}], 
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)
    # Initialize Amp.  Amp accepts either values or strings for the optional override arguments,
    # for convenient interoperation with argparse.
    [model, head], optimizer = amp.initialize([model, head], optimizer,
                                      opt_level=args.opt_level,
                                      keep_batchnorm_fp32=args.keep_batchnorm_fp32,
                                      loss_scale=args.loss_scale,
                                      min_loss_scale=-2.**45, 
                                      max_loss_scale=2.**45
                                      )
    # For distributed training, wrap the model with apex.parallel.DistributedDataParallel.
    # This must be done AFTER the call to amp.initialize.  If model = DDP(model) is called
    # before model, ... = amp.initialize(model, ...), the call to amp.initialize may alter
    # the types of model's parameters in a way that disrupts or destroys DDP's allreduce hooks.
    if args.distributed:
        # By default, apex.parallel.DistributedDataParallel overlaps communication with
        # computation in the backward pass.
        # model = DDP(model)
        # delay_allreduce delays all communication to the end of the backward pass.
        model = DDP(model, delay_allreduce=True)
        head = DDP(head, delay_allreduce=True)
    # define loss function (criterion) and optimizer
    criterion = FocalLoss().cuda()
    # Optionally resume from a checkpoint
    # if args.resume:
    #     # Use a local scope to avoid dangling references
    # CHECKPOINTS BACKBONE
    # Data loading code
    traindir = '/data'
    train_transform = transforms.Compose([ # refer to https://pytorch.org/docs/stable/torchvision/transforms.html for more build-in online data augmentation
        transforms.RandomHorizontalFlip(),
        # transforms.ToTensor(),
        # transforms.Normalize(mean = RGB_MEAN,
        #                      std = RGB_STD),
    ])
    all_dataset = Recognition_Dataset(traindir, train_transform)
    len_val = int(len(all_dataset)*0.1)
    len_train = len(all_dataset) - len_val
    train_dataset, val_dataset = torch.utils.data.random_split(all_dataset, [len_train, len_val])
    NUM_CLASS = len(train_dataset.dataset.classes)
    print("Number of Training Classes: {}, number train samples: {}, val samples: {}".format(NUM_CLASS, len_train, len_val))

    train_sampler = None
    val_sampler = None
    if args.distributed:
        print('start balance class sampler...')
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
        val_sampler  =  torch.utils.data.distributed.DistributedSampler(val_dataset)

    print('finish balance class sampler...')
    
    collate_fn = lambda b: fast_collate(b, memory_format)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler, collate_fn=collate_fn)
    

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True,
        sampler=val_sampler,
        collate_fn=collate_fn)

    if args.evaluate:
        validate(val_loader, model, head, criterion)
        return
    print('start train......')
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            train_sampler.set_epoch(epoch)

        # train for one epoch
        print('train in epoch...')
        train(train_loader, model, head, criterion, optimizer, epoch)

        # evaluate on validation set
        print('val')
        prec1 = validate(val_loader, model, head, criterion)

        # remember best prec@1 and save checkpoint
        if args.local_rank == 0:
            is_best = prec1 > best_prec1
            best_prec1 = max(prec1, best_prec1)
            save_checkpoint( model.state_dict(), head.state_dict(), epoch, is_best)

After starting training, as a rule, at the 3rd epoch, the output of the CNN model becomes NaN, while the image is normal at the input and such errors do not occur when starting training not through apex
Below I have printed out the gradients when an error occurs during training

Epoch: [2][5488/21528]	Time 0.356 (0.359)	Speed 1439.135 (1424.924)	Loss 30.7721385956 (29.7312)	Prec@1 0.000 (0.244)	Prec@5 0.391 (0.505)	Time 1952.462 sec 	 Epoch time 2.149 hours
Epoch: [2][5496/21528]	Time 0.356 (0.359)	Speed 1440.094 (1424.946)	Loss 30.6843338013 (29.7326)	Prec@1 0.000 (0.244)	Prec@5 0.000 (0.504)	Time 1954.005 sec 	 Epoch time 2.149 hours
Epoch: [2][5504/21528]	Time 0.354 (0.359)	Speed 1445.318 (1424.975)	Loss 30.5849666595 (29.7338)	Prec@1 0.195 (0.244)	Prec@5 0.391 (0.504)	Time 1949.777 sec 	 Epoch time 2.149 hours
grad sum
125.85698795318604 nan
545.2495787143707 nan
-12.65521240234375 nan
22.60831606388092 nan
3736.83203125 nan
17.25114631652832 nan
926.9287031888962 nan

In this case, the same error occurs when I am use optimization levels opt_level = O2
I could not run the model with a opt_level = O1, because the error in ArcFace forward is: expected scalar type float but found c10 :: Half. This error is gone, when I manually convert the data into ArcFace from float to Half, but in such case I get the error: Cuda out of memory for any batch size

Start by comparing the provided training pipeline with examples/imagenet and reproduce the NaN at about the third epoch under opt_level O2 or O3. Inspect the amp.initialize configuration, ArcFace forward path, and distributed model setup; done means identifying a reproducible Apex cause or a confirmed incompatibility with a documented mitigation.