Just an update on this, I guess the proper way is to free optimizer instead of the model, so it should be
optimizer.zero_grad(set_to_none=True), otherwisze the optimizer will keep the reference to the gradients, and thus they will keep place in memory.