In my case, memory usage was not changed at all while using the torch.cuda.amp
I don’t have a model (nn.Module) but optimizing a tensor, could that be the reason?
Note:
I have put them to nn.Parameter inside a model now and that didn’t change anything.