scale_loss: terminate called after throwing an instance of 'std::out_of_range'

Getting the following error when yield (loss.float())*loss_scale is called in scale_loss in handle.py:
terminate called after throwing an instance of 'std::out_of_range'
what(): vector::_M_range_check: __n (which is 335547169) >= this->size() (which is 3)
[1] 29602 abort (core dumped) python training.py

I installed the pip version of apex after cloning the latest repo and running: pip install -v --no-cache-dir ./

I initialize amp in the following way:
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
I added a print statement inside handle.py, just before the yield statement mentioned above.
These are the loss and loss_scale values right before the exception:
loss.float(): 0.07557302713394165, loss_scale: 65536.0
In this particular case, the exception came up 450 iterations into training. My training set has over 4000 samples.

I am training a 3d network for segmentation and due to GPU memory constraints, I pass the activation to GPU1 and the loss to the CPU, the remainder of the model runs on GPU0. This exception also comes up if I pass the loss to GPU1 (both activation and loss are computed in GPU1). The exception comes up randomly, sometimes towards the beginning of training, sometimes after several thousand iterations. I have also used this network architecture before, without problems.

Any tips on how to further debug this? Thanks!

Start with scale_loss in handle.py and the amp.initialize(..., opt_level='O2') setup; reproduce the intermittent failure using the reported split-device training configuration. Done means identifying the cause of the out_of_range failure and documenting or testing a fix.