[Tacotron2/PyTorch] zero_grad() unexpected keyword argument

[MikeyElmers]

Describe the bug I'm going through the Tacotron2 Quick Start Guide. I'm using the Tacotron 2 and WaveGlow PyTorch NGC container provided in step 3 in the quick start guide and I haven't made any changes to any of the scripts/files. When I run step 5 bash scripts/train_tacotron2.sh I get the following error

Traceback (most recent call last):
  File "train.py", line 557, in <module>
    main()
  File "train.py", line 509, in main
    model.zero_grad(set_to_none=True)
TypeError: zero_grad() got an unexpected keyword argument 'set_to_none'

Expected behavior The command should train Tacotron2.

Environment Please provide at least:

• Container version (e.g. pytorch:19.05-py3): NVIDIA-SMI 440.33.01
• GPUs in the system: (e.g. 8x Tesla V100-SXM2-16GB): 1x Tesla V100-PCIE-32GB
• CUDA driver version (e.g. 418.67): 11.0