Akjava commited on
Commit
2c05a28
Β·
verified Β·
1 Parent(s): 6b13f4e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -31,7 +31,7 @@ def init():
31
  model = AutoModelForCausalLM.from_pretrained(
32
  model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
33
  )
34
- text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device) #pipeline has not to(device)
35
 
36
  if next(model.parameters()).is_cuda:
37
  print("The model is on a GPU")
 
31
  model = AutoModelForCausalLM.from_pretrained(
32
  model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
33
  )
34
+ text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device,device=0 ) #pipeline has not to(device)
35
 
36
  if next(model.parameters()).is_cuda:
37
  print("The model is on a GPU")