What is the correct way to edit parameter?

The code ( https://gist.github.com/cswhjiang/be475ef9a3a7d1f781830ebfb7970719 ) failed on pytorch 0.2. What is the correct way to edit parameter?

python3.6 -u weight_drop.py
Applying weight drop of 0 to weight_hh_l0
[['weight_ih_l0', 'bias_ih_l0', 'bias_hh_l0', 'weight_hh_l0_raw']]
odict_keys(['weight_ih_l0', 'bias_ih_l0', 'bias_hh_l0', 'weight_hh_l0_raw'])
Traceback (most recent call last):
  File "weight_drop.py", line 44, in <module>
    a.cuda()
  File "/data1/XXX/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 147, in cuda
    return self._apply(lambda t: t.cuda(device_id))
  File "/data1/XXX/.local/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 116, in _apply
    self.flatten_parameters()
  File "/data1/XXX/.local/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 107, in flatten_parameters
    rnn._copyParams(all_weights, params)
  File "/data1/XXX/.local/lib/python3.6/site-packages/torch/backends/cudnn/rnn.py", line 187, in _copyParams
    param_to.copy_(param_from, broadcast=False)
RuntimeError: invalid argument 2: sizes do not match at /pytorch/torch/lib/THC/THCTensorCopy.cu:31