Jeff Myers II commited on
Commit
a4a81f7
·
1 Parent(s): d1c2a25

Removed the 'device' arg from pipeline

Browse files
Files changed (1) hide show
  1. Gemma.py +2 -1
Gemma.py CHANGED
@@ -14,7 +14,8 @@ class GemmaLLM:
14
  model_id = "google/gemma-3-4b-it"
15
  model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, pad_token_id=0)
16
  tokenizer = AutoTokenizer.from_pretrained(model_id)
17
- self.model = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.bfloat16, device="auto")
 
18
 
19
  @spaces.GPU
20
  def generate(self, message) -> str:
 
14
  model_id = "google/gemma-3-4b-it"
15
  model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, pad_token_id=0)
16
  tokenizer = AutoTokenizer.from_pretrained(model_id)
17
+
18
+ self.model = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.bfloat16)
19
 
20
  @spaces.GPU
21
  def generate(self, message) -> str: