src/utils.py

from transformers import ViTForImageClassification
import os

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))


def load_model():
    return ViTForImageClassification.from_pretrained(ROOT_DIR)