Hello,

I am using a U-Net augmentation (specifically: https://github.com/juntang-zhuang/LadderNet) to perform segmentation of hands. To be specific, I am classifying each pixel of an image to one of five classes (no hand, my right hand, my left hand, your right hand, your left hand.)

This requires my prob shape (in fisher.py: 446) to be [batch_size, img_h, img_w, n_classes] -> [8,32,32,5] (snippet below)

def __fisher_exact(loss_and_backward, model, probs):
    _, n_classes = probs.shape  

Because of this dimensionality, this line of code fails. I assume it's because it expects the probs.shape tuple to be (img_as_tensor, label_as_int) per the CIFAR example: https://github.com/AlexImmer/Laplace/blob/main/examples/calibration_example.py) where the CIFAR dataset object returns a a tuple of (#_examples, (img_as_tensor: [3,32,32], label_as_int: 0,1,2,3,etc.).

I can always reshape my training data to be a tuple of that format, but because I am classifying by pixel, my associated label would not be a single integer. It would have to be in the same shape as the image tensor [32,32].

So I'm asking this question to the community to see if anyone has attempted this kind of segmentation task using the laplace-torch package before I try to force a solution within the asdfghjkl/fisher.py file.