Thank for that tip. That gets rid of the keyerror but now I receive the following error (stack trace below):
Traceback (most recent call last):
  File "run_experiments.py", line 118, in <module>
    train_model(model, train_loader, nb_batches, optimizer, criterion, **vars(args))
  File "run_experiments.py", line 47, in train_model
    loss.backward()
  File "/if1/ab3cb/miniconda3/envs/fast_torch/lib/python3.5/site-packages/torch/autograd/variable.py", line 146, in backward
    self._execution_engine.run_backward((self,), (gradient,), retain_variables)
  File "/if1/ab3cb/miniconda3/envs/fast_torch/lib/python3.5/site-packages/torch/nn/_functions/linear.py", line 22, in backward
    grad_input = torch.mm(grad_output, weight)
TypeError: torch.mm received an invalid combination of arguments - got (torch.FloatTensor, torch.cuda.FloatTensor), but expected one of:
 * (torch.SparseFloatTensor mat1, torch.FloatTensor mat2)
      didn't match because some of the arguments have invalid types: (torch.FloatTensor, torch.cuda.FloatTensor)
 * (torch.FloatTensor source, torch.FloatTensor mat2)
      didn't match because some of the arguments have invalid types: (torch.FloatTensor, torch.cuda.FloatTensor)

I dont understand why the loss.backward() operation expects anything to be a non-cuda tensor when I have converted my model and associated parameters to cuda tensors. Am I missing something here?
PS: I have updated the gists to reflect the new code.