CUDA error while passing hidden states to convLSTM - conv-neural-network

Question

I'm getting the following error, here's the traceback:
Traceback (most recent call last):
File "train.py", line 136, in <module>
train(epoch)
File "train.py", line 112, in train
output = model(data, hc) # Get outputs of LSTM
File "/home/ama1128/.conda/envs/matrix/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
result = self.forward(*input, **kwargs)
File "/scratch/ama1128/convLSTM/model.py", line 118, in forward
hc = self.cell_list[t](input=x[0], prev_state=[h_0, c_0])
File "/home/ama1128/.conda/envs/matrix/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
result = self.forward(*input, **kwargs)
File "/scratch/ama1128/convLSTM/model.py", line 50, in forward
combined = torch.cat((input, h_prev), dim=1) # concatenate along channel axis
RuntimeError: Expected object of backend CUDA but got backend CPU for sequence element 1 in sequence argument at position #1 'tensors'
Training Loop:
def train(epoch):
model.train()
hc = model.init_hidden(batch_size=1)
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
target = data[1:] # Set target, images 2 to 18
if gpu:
data = data.cuda()
target = target.cuda()
output = model(data, hc) # outputs of LSTM
loss = criterion(output, target)
loss.backward()
optimizer.step()
convLSTMCell forward:
def forward(self, input, prev_state):
h_prev, c_prev = prev_state
combined = torch.cat((input, h_prev), dim=1) # concatenate along channel axis
combined_conv = self.conv(combined)
cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
i = torch.sigmoid(cc_i)
f = torch.sigmoid(cc_f)
o = torch.sigmoid(cc_o)
g = torch.tanh(cc_g)
c_cur = f * c_prev + i * g
h_cur = o * torch.tanh(c_cur)
return h_cur, c_cur
convLSTM Class:
Initialize states:
def init_hidden(self, batch_size):
hidden = torch.zeros(batch_size, self.hidden_dim[0], self.height, self.width)
cell = torch.zeros(batch_size, self.hidden_dim[0], self.height, self.width)
if gpu:
hidden.cuda()
cell.cuda()
return hidden, cell
convLSTM forward:
def forward(self, x, hc):
outputs = []
states = []
x = torch.unsqueeze(x, 1) # change shape from (18,3,128,128) to (18,1,3,128,128)
h_0, c_0 = hc
for t in range(0, self.num_layers):
if t==0:
hc = self.cell_list[t](input=x[0], prev_state=[h_0, c_0])
states.append(hc)
hidden, cell = hc
outputs.append(hidden.view(3,128,128))
else:
h, c = states[t-1] # unpack previous states
if gpu:
h.cuda()
c.cuda()
hc = self.cell_list[t](input=x[t], prev_state=[h,c]) # current states
states.append(hc) # store current states for next cell
hidden, cell = hc # unpack current states
outputs.append(hidden.view(3,128,128))
outputs is reformatted using torch.stack before passing it to the loss function.
I've tried to cast the hidden states to gpu wherever I can and the error still persists. What's going on?