I declared helper function and not getting the error anymore:

def compute_loss(criterion, outputs, targets):
“”"
Helper function to compute the loss. Since this is a pixelwise
prediction task we need to reshape the output and ground truth
tensors into a 2D tensor before passing it in to the loss criteron.
Args:
  criterion: pytorch loss criterion
  outputs (pytorch tensor): predicted labels from the model
  labels (pytorch tensor): ground truth labels
Returns:
  pytorch tensor for loss
"""

loss_out = outputs
loss_targs = targets
return criterion(loss_out, loss_targs)


However, the model seems to be stuck 

Exception reporting mode: Verbose
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/container.py:92: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
input = module(input)
tensor([[0.2053, 0.1828, 0.1825, 0.2460, 0.1834]],
grad_fn=)
tensor([0, 0])