The code should work for a segmentation model output of [N, nb_classes, H, W] and the corresponding target of [N, H, W] as this code snippet shows:
N, H, W = 2, 224, 224
nb_classes = 10
output = torch.randn(N, nb_classes, H, W)
preds = torch.argmax(output, 1)

labels = torch.randint(0, nb_classes, (N, H, W))

print(torch.sum(preds == labels).float() / labels.nelement())
> tensor(0.989)