Gradient Overflow & Prediction FP16 Input FP32

I have tried to use APEX's AMP O1 I have the following problems:

1. My network input is FP32 but the prediction is FP16 which results in PyTorch errors since the types are not the same
2. If I try to fix 1. by casting the prediction to FP32 via tensor.float() I start running into Gradient Overflows. They happen nearly continuously, I have 1 - 2 batches that do not overflow here and there but very rarely. Also, it leads to non-finite values (inf,-inf,nan) in my predictions.

If I change to us O2 I have yet another mixed FP32 FP16 operation error in my code. The documentation does not help at all for this, so my question is, can I tell AMP to keep a certain tensor to FP32 or keep all the tensors in the loss to FP32?

Start with the reported O1/O2 mixed-precision behavior and the linked APEX advanced documentation section on forcing particular layers or functions to a desired type. Reproduce the FP16 prediction, FP32 cast, and gradient-overflow symptoms if the original network is available; done means documenting or correcting the supported way to keep the relevant tensor or loss computation in FP32.