Update models/model_loader.py

models/model_loader.py

@@ -3,6 +3,8 @@ import streamlit as st
 
 CACHE_DIR = "/tmp/pretrained_models"
 
+import torch
+import gc
 
 from speechbrain.pretrained.interfaces import foreign_class
 from faster_whisper import WhisperModel
@@ -34,3 +36,8 @@ def load_accent_model():
 @st.cache_resource(show_spinner="Loading Whisper...")
 def load_whisper():
     return WhisperModel("tiny", device="cpu", compute_type="int8_float32")
+
+def unload_model(model):
+    del model
+    torch.cuda.empty_cache()
+    gc.collect()