commit bd054e7cef00da543b80844aeb7ae0ac78c91675
author meijieru <meijieru@gmail.com> Tue Aug 13 12:14:03 2019 -0700
committer Facebook Github Bot <facebook-github-bot@users.noreply.github.com> Tue Aug 13 12:18:31 2019 -0700
tree 8ccf369b6ad9c6db020c881746bc4e23a74a6fca
parent 5ae909b443570ae89b3716185cb9d26e4081a619

reduce memory usage for centered rmsprop (#24170)

Summary:
Reduce gpu memory usage by using in-place operation
Pull Request resolved: https://github.com/pytorch/pytorch/pull/24170

Differential Revision: D16784495

Pulled By: vincentqb

fbshipit-source-id: 03820cdc9a3952b95b9af0f87d3a9bb0f21e9b4d

torch/optim/rmsprop.py

diff --git a/torch/optim/rmsprop.py b/torch/optim/rmsprop.py
index cbb28f7..caa23af 100644
--- a/torch/optim/rmsprop.py
+++ b/torch/optim/rmsprop.py
@@ -88,7 +88,7 @@
                 if group['centered']:
                     grad_avg = state['grad_avg']
                     grad_avg.mul_(alpha).add_(1 - alpha, grad)
-                    avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt().add_(group['eps'])
+                    avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt_().add_(group['eps'])
                 else:
                     avg = square_avg.sqrt().add_(group['eps'])