@putopavel Sorry, my mistake. This is how its done
config = AutoConfig.from_pretrained(model_name, num_lables=3)
model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)

Also would be nice if you only post the relevant code snippet and post colab link for full code. Kinda hard to read when there’s so much code