Implement NLLLossNd

[zou3519]

Needed for #3556

I'm not sure this is the best way to implement because the .contiguous() calls might be slow.

One alternative way to implement this is to copy and modify gather. Without any of the extra keyword modifiers, with reduce=False, the following is equivalent to NLLLossNd:

def nlllossNd(input, target):
    target = target.unsqueeze(1)
    out = torch.gather(input, 1, target)
    return out.squeeze(1)

I tried benchmarking this against what I have right now (this diff that uses .contiguous() calls and NLLLoss2d) and using gather is around 2x slower, even for non-contiguous inputs, so I went with this approach.

Test Plan

Unit tests for NLLLossNd with a NLLLossNd reference function

[zou3519]

@pytorchbot retest this please

[pietern]

There was some CI maintenance happening this morning -- retriggering build.

@pytorchbot retest this please