Spaces:

Ruurd
/

radiolm

Running on Zero

Ruurd commited on 30 days ago

Commit

f6e6d69

1 Parent(s): b6b7c74

Put model on GPU

Files changed (1) hide show

app.py CHANGED Viewed

@@ -19,7 +19,7 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     current_model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,
-        device_map="auto",
         use_auth_token=token
     )

     current_model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,
+        device_map="cuda",
         use_auth_token=token
     )