Spaces:
Paused
Paused
Commit
·
6de173e
1
Parent(s):
5d04704
Update app_v4.py
Browse files
app_v4.py
CHANGED
@@ -33,7 +33,7 @@ if device == "cuda:0":
|
|
33 |
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=False)
|
34 |
|
35 |
# Attempt to load the model, catch any OOM errors
|
36 |
-
@st.
|
37 |
def load_gptq_model():
|
38 |
model = AutoGPTQForCausalLM.from_quantized(
|
39 |
pretrained_model_dir,
|
@@ -45,10 +45,8 @@ def load_gptq_model():
|
|
45 |
return model
|
46 |
|
47 |
# Attempt to load the model, catch any OOM errors
|
48 |
-
model_loaded = False
|
49 |
try:
|
50 |
model = load_gptq_model()
|
51 |
-
model_loaded = True
|
52 |
except RuntimeError as e:
|
53 |
if 'CUDA out of memory' in str(e):
|
54 |
st.error("CUDA out of memory while loading the model. Try reducing the model size or restarting the app.")
|
|
|
33 |
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=False)
|
34 |
|
35 |
# Attempt to load the model, catch any OOM errors
|
36 |
+
@st.cache_resource
|
37 |
def load_gptq_model():
|
38 |
model = AutoGPTQForCausalLM.from_quantized(
|
39 |
pretrained_model_dir,
|
|
|
45 |
return model
|
46 |
|
47 |
# Attempt to load the model, catch any OOM errors
|
|
|
48 |
try:
|
49 |
model = load_gptq_model()
|
|
|
50 |
except RuntimeError as e:
|
51 |
if 'CUDA out of memory' in str(e):
|
52 |
st.error("CUDA out of memory while loading the model. Try reducing the model size or restarting the app.")
|