[solved] When running on gpu cuda type error for loss function

Hi

On CPU the model I implemented worked ok, but when converting it to cuda, it gives the following error

Traceback (most recent call last):
File “train_main.py”, line 281, in
main()
File “train_main.py”, line 270, in main
trn_loss, trn_acc = _trn_epoch(model, epoch, batchid)
File “train_main.py”, line 235, in _trn_epoch
loss.backward()
File “/home/hgodhia/miniconda2/envs/anlp/lib/python2.7/site-packages/torch/autograd/variable.py”, line 146, in backward
self._execution_engine.run_backward((self,), (gradient,), retain_variables)
File “/home/hgodhia/miniconda2/envs/anlp/lib/python2.7/site-packages/torch/nn/_functions/thnn/auto.py”, line 175, in backward
update_grad_input_fn(self._backend.library_state, input, grad_output, grad_input, *gi_args)
TypeError: CudaSoftMax_updateGradInput received an invalid combination of arguments - got (int, torch.cuda.FloatTensor, torch.FloatTensor, torch.cuda.FloatTensor, torch.cuda.FloatTensor), but expected (int state, torch.cuda.FloatTensor input, torch.cuda.FloatTensor gradOutput, torch.cuda.FloatTensor gradInput, torch.cuda.FloatTensor output)

I have checked thoroughly through the code and where ever I used tensors I changed them to cuda tensors like

    if torch.cuda.is_available():
        zero_t = zero_t.cuda(0)
        end_idxs_flat = end_idxs_flat.cuda(0)

including the embedding layers etc
rest assured the input to loss
loss = loss_function(scores, a)
are cuda tensor variables
loss_function = nn.NLLLoss()

if torch.cuda.is_available():
loss_function = loss_function.cuda(0)

Silly error, RESOLVED.
Thanks