Out of memory problem in amp.initialize

in this case, I have loaded the model to cuda:2, but when I run amp.initialize, I find ' self._overflow_buf' is loaded to default cuda device which is cuda:0. But cuda:0 memory has been exhausted by another process.

Here is the error message.

----> model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

/data/scratch/owaski/anaconda3/lib/python3.7/site-packages/apex/amp/frontend.py in initialize(models, optimizers, enabled, opt_level, cast_model_type, patch_torch_functions, keep_batchnorm_fp32, master_weights, loss_scale, cast_model_outputs, num_losses, verbosity, min_loss_scale, max_loss_scale)
    356         maybe_print("{:22} : {}".format(k, v), True)
    357
--> 358     return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)
    359
    360

/data/scratch/owaski/anaconda3/lib/python3.7/site-packages/apex/amp/_initialize.py in _initialize(models, optimizers, properties, num_losses, cast_model_outputs)
    228         _amp_state.loss_scalers.append(LossScaler(properties.loss_scale,
    229                                                   min_loss_scale=_amp_state.min_loss_scale,
--> 230                                                   max_loss_scale=_amp_state.max_loss_scale))
    231
    232     if properties.patch_torch_functions:

/data/scratch/owaski/anaconda3/lib/python3.7/site-packages/apex/amp/scaler.py in __init__(self, loss_scale, init_scale, scale_factor, scale_window, min_loss_scale, max_loss_scale)
     54         self._unskipped = 0
     55         self._has_overflow = False
---> 56         self._overflow_buf = torch.cuda.IntTensor([0])
     57         if multi_tensor_applier.available:
     58             import amp_C

RuntimeError: CUDA error: out of memory

Start at apex/amp/scaler.py in LossScaler.init, where _overflow_buf is created, and follow its call from apex/amp/_initialize.py during amp.initialize. Reproduce with the model on cuda:2 while cuda:0 is full; done when initialization no longer allocates the overflow buffer on the default device and avoids the reported out-of-memory error.