train.py

from transformers import AutoModelForSequenceClassification, AutoTokenizer

def main():
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    model = AutoModelForSequenceClassification.from_pretrained(
        "bert-base-uncased", num_labels=3
    )

    text = "pick up the object"
    inputs = tokenizer(text, return_tensors="pt")

    outputs = model(**inputs)
    print(outputs.logits)

if __name__ == "__main__":
    main()