Issue description
spectral_norm used in nn.Linear is okay. But when it’s used in nn.RNN,there will be a RuntimeError while running model = network().cuda()
Traceback (most recent call last):
File “/DATA/119/yhe/soft/anaconda3/envs/env1/lib/python3.6/site-packages/IPython/core/interactiveshell.py”, line 3296, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File “”, line 1, in
runfile(’/DATA/119/yhe/code/crit/mnist/example.py’, wdir=’/DATA/119/yhe/code/crit/mnist’)
File “/home/yhe/.pycharm_helpers/pydev/_pydev_bundle/pydev_umd.py”, line 197, in runfile
pydev_imports.execfile(filename, global_vars, local_vars) # execute the script
File “/home/yhe/.pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py”, line 18, in execfile
exec(compile(contents+"\n", file, ‘exec’), glob, loc)
File “/DATA/119/yhe/code/crit/mnist/example.py”, line 17, in
model = network().cuda()
File “/DATA/119/yhe/soft/anaconda3/envs/env1/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 265, in cuda
return self._apply(lambda t: t.cuda(device))
File “/DATA/119/yhe/soft/anaconda3/envs/env1/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 193, in _apply
module._apply(fn)
File “/DATA/119/yhe/soft/anaconda3/envs/env1/lib/python3.6/site-packages/torch/nn/modules/rnn.py”, line 127, in _apply
self.flatten_parameters()
File “/DATA/119/yhe/soft/anaconda3/envs/env1/lib/python3.6/site-packages/torch/nn/modules/rnn.py”, line 123, in flatten_parameters
self.batch_first, bool(self.bidirectional))
RuntimeError: param_from.type() == param_to.type() ASSERT FAILED at /pytorch/aten/src/ATen/native/cudnn/RNN.cpp:541, please report a bug to PyTorch. parameter types mismatch
Code
import torch
import torch.nn as nn
import torch.nn.functional as F
class network (nn.Module):
def __init__(self):
super(network, self).__init__()
self.sn_rnn = nn.utils.spectral_norm(nn.RNN(20,10,1),name='weight_hh_l0')
# self.fc = nn.utils.spectral_norm(nn.Linear(20,10), name='weight')
def forward (self,x):
# x = F.tanh(self.fc(x))
out,_ = self.sn_rnn(x)
return F.log_softmax(x, dim=1)
if __name__ == '__main__':
x = torch.randn(2, 10, 20).cuda()
model = network().cuda()
out= model(x)
print('end')
Environment
- PyTorch Version (e.g., 1.0): 1.1.0
- Python version: 3.6.8
- CUDA/cuDNN version: cuda9.0