Solved by ptrblck in post #2 
              
              
                apex.amp is deprecated so use the native mixed-precision training utility via torch.cuda.amp. You can find the examples here.