When I remove is_inference from:

class NewSegmentationModule(SegmentationModule):
def forward(self, image, label, is_inference=None):
segSize = (label.shape[-2], label.shape[-1])
pred = self.decoder(self.encoder(image, return_feature_maps=True), segSize=segSize, is_inference=is_inference)
return pred

The error is gone but I get another error: