The permutations assume 4-dimensional tensors.
Here comes the first difference to Keras/TF: In PyTorch these will be Batch, Channel/Class, Height, Width, wit the channel containing the class label (in TF it’s BHWC, as pointed out in the comment you linked).
So what you want is that TP FN and FP sum over B, H and W (you could do that by doing torch.sum(label * pred, dim=(0, 2, 3)) apparently, so you would get a vector with just the class dimension, i.e. a score per class.
Note that this convention is different to the usual averaging over batches which may or may not affect what you want to do when changing batch sizes (depends on optimizer etc.).
Best regards
Thomas