Update app.py

app.py

@@ -52,12 +52,15 @@ def reset_model(model_name, *extra_components, with_extra_components=True):
     global_state.layers_format = model_args.pop('layers_format')
     tokenizer_path = model_args.pop('tokenizer') if 'tokenizer' in model_args else model_path
     use_ctransformers = model_args.pop('ctransformers', False)
+    dont_cuda = model_args.pop('dont_cuda', False)
     AutoModelClass = CAutoModelForCausalLM if use_ctransformers else AutoModelForCausalLM
     
     # get model
     global_state.model, global_state.tokenizer, global_state.hidden_states = None, None, None
     gc.collect()
-    global_state.model = AutoModelClass.from_pretrained(model_path, **model_args).to('cuda')
+    global_state.model = AutoModelClass.from_pretrained(model_path, **model_args)
+    if not dont_cuda:
+        global_state.model.to('cuda')
     global_state.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, token=os.environ['hf_token'])
     gc.collect()
     if with_extra_components: