TinNguyenqn:

optimizer = optim.SGD(model.parameters(), lr=args.lr)


You are passing the model parameters into the optimizer. So, the model parameters would be optimized (using SGD) and not the auto_encoder_model parameters. Unless there are pieces of code which are missing here.