from transformers import AutoFeatureExtractor, AutoModelForImageClassification import torch from datasets import load_dataset model = AutoModelForImageClassification.from_pretrained(".", trust_remote_code=True) dataset = load_dataset("huggingface/cats-image") image = dataset["test"]["image"][0] feature_extractor = AutoFeatureExtractor.from_pretrained(".") inputs = feature_extractor(image, return_tensors="pt") with torch.no_grad(): logits = model(**inputs).logits # model predicts one of the 1000 ImageNet classes predicted_label = logits.argmax(-1).item() print(model.config.id2label[predicted_label])