Thanks so much @ptrblck. I have changed my code using autocast for torch.cuda.amp. However I remark that before with
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")

I can use a batch_size of 16, however using the autocast, I could not use more than 8 !!! It seems some operations are not casted to float16 as in apex! I think the autocast implement the level 01 of apex