Using the following model:
class Model(nn.Module):
def __init__(self, n_classes, num_layers, hidden_size):
super(Model, self).__init__()
self.blstm = nn.LSTM(input_size=96, hidden_size=hidden_size, num_layers=num_layers, dropout=0, bidirectional=True)
self.fc = nn.Linear(hidden_size * 2, n_classes)
self.softmax = nn.LogSoftmax(dim=2)
def forward(self, x):
# x of shape bs, 1, h, w
x = x.squeeze(1)
# x of shape bs, h, w
x = x.permute(2, 0, 1)
# x of shape w, bs, h
x, _ = self.blstm(x, None)
# x of shape w, bs, 512
x = self.fc(x)
# x of shape w, bs, n_classes
x = self.softmax(x)
return x
I get the following error:
x, _ = self.blstm(x, None)
File "C:\Users\Mohammed\.conda\envs\dl_env\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "C:\Users\Mohammed\.conda\envs\dl_env\lib\site-packages\torch\nn\modules\rnn.py", line 581, in forward
result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED
This error happens only on GPU, when I run my model on CPU everything works fine.