Thanks a lot for looking into this. I can confirm that the realistic setting with different xs leads to more memory use for AMP=False as expected. I am considering this issue resolved 
PS: About manual scaler values. The issue is that scaler.step() method requires an optimizer as input. In my case, I cannot use Pytorch optimizers because I need weights at step t to be a function of the weights at step t-1, t-2, ..., 0, in order to backprop through the whole inner loop. Furthermore, AFAIK GradScaler simply decreases the scaler value when nans occur, but I grid-searched a large range of scaler values and always get nans in my code, so I don’t think GradScaler would help in my case. Unless nans can occur for reasons I haven’t considered I’m guessing some computational graphs are just too deep/chaotic for AMP.