finetuning from FP32 model

Hi,
Is that a feasible way to finetuning from a fp32 model in amp mode? I've tried this but got loss Nan.
My code is written in this way:

  net = model.SE_LResNet100E_IR()
  # convert to sync_bn
  net = apex.parallel.convert_syncbn_model(net)
  net.to(device)
  optimizer = torch.optim.SGD(net.parameters(),
                              lr=lr, weight_decay=weight_decay, momentum=momentum)
  if use_fp16:
    master_print("Initialize AMP...")
    net, optimizer = amp.initialize(net, optimizer,
                                    opt_level="O2")
  if is_dist:
    net = nn.parallel.DistributedDataParallel(net, device_ids=device_ids, output_device=rank)
  if pretrain:
    master_print("Loading from pretrained fp32 weights: ", pretrain)
    pretrain_state = torch.load(pretrain, map_location='cpu')
    net.load_state_dict(pretrain_state['state_dict'])

When use_fp16 is disabled, the loss value keeps normal and converges through the training process.

Start with the reported sequence around amp.initialize(..., opt_level="O2") and the later pretrained net.load_state_dict call. Reproduce the NaN loss while loading FP32 weights, compare it with the non-AMP path, and determine whether this loading flow is supported or needs a documented correction.