Update app_v4.py

app_v4.py

@@ -15,9 +15,9 @@ def get_gpu_memory():
         memory_info = [{"free": "N/A", "total": "N/A"}]
     return memory_info
 
-# Display GPU memory information
-gpu_memory = get_gpu_memory()
-st.write(f"GPU Memory Info: {gpu_memory}")
+# Display GPU memory information before loading the model
+gpu_memory_before = get_gpu_memory()
+st.write(f"GPU Memory Info before loading the model: {gpu_memory_before}")
 
 # Define pretrained model directory
 pretrained_model_dir = "FPHam/Jackson_The_Formalizer_V2_13b_GPTQ"
@@ -43,11 +43,15 @@ try:
     )
 except RuntimeError as e:
     if 'CUDA out of memory' in str(e):
-        st.error("CUDA out of memory before toking generation. Try reducing the model size or input length.")
+        st.error("CUDA out of memory while loading the model. Try reducing the model size or input length.")
         st.stop()
     else:
         raise e
 
+# Display GPU memory information after loading the model
+gpu_memory_after = get_gpu_memory()
+st.write(f"GPU Memory Info after loading the model: {gpu_memory_after}")
+
 # User input for the model
 user_input = st.text_input("Input a phrase")