Update app.py
Browse files
app.py
CHANGED
@@ -72,8 +72,10 @@ def reset_model(model_name, load_on_gpu, *extra_components, reset_sentence_trans
|
|
72 |
if reset_sentence_transformer:
|
73 |
global_state.sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')
|
74 |
gc.collect()
|
75 |
-
if
|
76 |
global_state.model.to('cuda')
|
|
|
|
|
77 |
global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
|
78 |
gc.collect()
|
79 |
if with_extra_components:
|
|
|
72 |
if reset_sentence_transformer:
|
73 |
global_state.sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')
|
74 |
gc.collect()
|
75 |
+
if not dont_cuda:
|
76 |
global_state.model.to('cuda')
|
77 |
+
if load_on_gpu:
|
78 |
+
global_state.model.to('cpu')
|
79 |
global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
|
80 |
gc.collect()
|
81 |
if with_extra_components:
|