Fine Tuning OneFormer for Semantic Segmentation, how to change model's class number

[eBeyzaG]

I am fine tuning OneFormer for semantic segmentation and I have 5 classes including background. While training, I am also calculating mIOU so I extract segmentation maps. But as the model config has 150 classes, it outputs class numbers that do not exist in my dataset.
I also tried changing num_labels parameter as shown below

model=OneFormerForUniversalSegmentation.from_pretrained(
        "shilabs/oneformer_ade20k_swin_tiny",
        num_labels=5,
        ignore_mismatched_sizes=True)

But model gives error

446 # [batch_size, hidden_dim]
447 image_queries = nn.functional.normalize(image_queries.flatten(1), dim=-1)
--> 448 text_queries = nn.functional.normalize(text_queries.flatten(1), dim=-1)
450 logit_scale = torch.clamp(self.logit_scale.exp(), max=100)
452 logits_per_text = torch.matmul(text_queries, image_queries.t()) * logit_scale

AttributeError: 'NoneType' object has no attribute 'flatten'

How can I change class count while fine tuning?

[BhavanaMP]

@eBeyzaG I guess you have to add your classes id2label and label2id to the model while instantiating

something like this...
model = AutoModelForUniversalSegmentation.from_pretrained(
model_ckpt, is_training=is_train, id2label=id2label,
label2id=label2id, ignore_mismatched_sizes=True
)

[eBeyzaG]

Thank you for answering. While trying your answer I realized the real problem, I didn't specify is_training=True while instantiating the model causing the text queries to not be calculated. The error was caused by this.