src/pipeline/utils.py

import os
import logging
import torch
import torch.nn as nn
from path_config import MODEL_DIR

def save_model(model, model_name):
    """
    Save the trained model
    """
    os.makedirs(MODEL_DIR, exist_ok=True)

    model_path = os.path.join(MODEL_DIR, model_name)
    torch.save(model.state_dict(), model_path)
    logging.info(f"Model saved at {model_path}")
    print("Saved successfully!")