move model

src/train.py

@@ -5,7 +5,7 @@ from transformers import (
     Trainer,
 )
 from datasets import load_dataset
-from .utils import MODEL_DIR
+from .utils import ROOT_DIR
 
 
 def train():
@@ -67,5 +67,5 @@ def train():
     trainer.train()
 
     # Save model and processor
-    model.save_pretrained(MODEL_DIR)
-    processor.save_pretrained(MODEL_DIR)
+    model.save_pretrained(ROOT_DIR)
+    processor.save_pretrained(ROOT_DIR)

src/utils.py

@@ -2,8 +2,7 @@ from transformers import ViTForImageClassification
 import os
 
 ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-MODEL_DIR = os.path.join(ROOT_DIR, "model")
 
 
 def load_model():
-    return ViTForImageClassification.from_pretrained(MODEL_DIR)
+    return ViTForImageClassification.from_pretrained(ROOT_DIR)