Update to work with pytorch 1.11.0

[aleloi]

Previously running this on pytorch 1.11.0 failed with RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Int'. After converting training labels to long (which is not for some reason a requirement), I ran in the next error RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.. I fixed both and the accuracy/loss is computed normally.

[YangZhangMizzou]

Thank you. It worked on my code.