Thanks for the response! Now the torch amp can run correctly. Another quick question: how can I switch between Apex amp O0 and O1/O2/O3 during training time? Is model checkpointing and reloading the only way to accomplish this, since the model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level) is set before the training begins?