def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll