TAO toolkit yolo_v4 training NAN loss

Hello,

I keep getting NAN losses early in my training with TAO 4.0.1 both in the provided jupyter notebook and in my personal environment with my own dataset.

After some investigation, I have realised that setting --use_amp causes the loss to NAN.

Just wanted to flag this!

Best,
Benedetta

Thanks for the info.