Fix "RuntimeError: CUDA Error: out of memory"

[icekang]

Problem

According to this issue
I also forked this repo and tried running on my Colab project. The same problem arose: RuntimeError: CUDA Error: out of memory.

Solution

As far as my knowledge, the problem happened because some section did not require grad yet still did anyway. Thus, with no_grad() should be presented.

[ghost]

@icekang Thank you!! Changes worked fine and saved me debugging time.

[icekang]

you are welcome :)
if it works, please accept my pull request @arnoweng