Solved by vfdev-5 in post #2 
              
              
                There are several repositories where it was already implemented (probably in the correct way): 

https://github.com/pytorch/vision/blob/master/references/segmentation/utils.py#L95
https://github.com/pytorch/ignite/blob/master/ignite/metrics/confusion_matrix.py#L129

In your case, as far as I underst…