SegFormer fine-tuned on personal dataset problem on loss computation

Hi I have a problem with the fine tuning of SegFormer model.
I start from the code here and when I try to execute on my own pc it works.
I modify the dataset class to open my own dataset and when i pass the pixel_values and the label to the model it give me a problem about Unable to get repr for class torch.Tensor after the model has taken the image as input and analyzed it.
I share with you the code that I have changed for dataset loading:

class SemanticSegmentationDataset(Dataset):
    """Image (semantic) segmentation dataset."""

    def __init__(self, dataset_info:dict, feature_extractor, dataset_type = 'train'):

        self.feature_extractor = feature_extractor

        if dataset_type == 'train':
            self.images = dataset_info['train']['pixel_values']
            self.mask = dataset_info['train']['label']
        elif dataset_type == 'val':
            self.images = dataset_info['val']['pixel_values']
            self.mask = dataset_info['val']['label']
        elif dataset_type == 'test':
            self.images = dataset_info['test']['pixel_values']
            self.mask = dataset_info['test']['label']
        else:
            raise NotImplementedError('Type of dataset not managed')


    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):

        image = imread(self.images[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        segmentation_map = imread(self.mask[idx], cv2.IMREAD_UNCHANGED)

        # randomly crop + pad both image and segmentation map to same size
        encoded_inputs = self.feature_extractor(image, segmentation_map, return_tensors="pt")

        for k, v in encoded_inputs.items():
            encoded_inputs[k].squeeze_()  # remove batch dimension

        return encoded_inputs

I have also tried with the use of skimage for opening the image and also with the use of PIL. More over I have used also the example that use the dataset.Dataset class but everytime the result is the same.

Can you help me please?