Thanks Patrick. You were correct.
Target value was out of bounds and running on CPU gave a much better stack trace.
This is a custom dataset that I created on my own. Look at the stack trace below, if you could tell me if the error is with the data or the model architecture that would be amazing:
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-12-de1f45d3ba4c> in <module>
     22         idmask_pred,umask_pred,vmask_pred = correspondence_block(image)
     23         # calculate the batch loss
---> 24         loss_id = criterion_id(idmask_pred, idmask)
     25         loss_u = criterion_u(umask_pred, umask)
     26         loss_v = criterion_v(vmask_pred, vmask)

/opt/venv/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

/opt/venv/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    930     def forward(self, input, target):
    931         return F.cross_entropy(input, target, weight=self.weight,
--> 932                                ignore_index=self.ignore_index, reduction=self.reduction)
    933 
    934 

/opt/venv/lib/python3.7/site-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
   2315     if size_average is not None or reduce is not None:
   2316         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2317     return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
   2318 
   2319 

/opt/venv/lib/python3.7/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2115         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2116     elif dim == 4:
-> 2117         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2118     else:
   2119         # dim == 3 or dim > 4