Nick088 commited on
Commit
31328c7
1 Parent(s): c1a0ecb

Fixed expected all tensors to be on the same device

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -12,6 +12,8 @@ else:
12
  tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
13
  model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", torch_dtype=torch.float16)
14
 
 
 
15
 
16
  def generate(
17
  prompt, history, max_new_tokens=512, repetition_penalty=1.2, temperature=0.5, top_p=1, top_k=1, seed=42
 
12
  tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
13
  model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", torch_dtype=torch.float16)
14
 
15
+ model.to(device)
16
+
17
 
18
  def generate(
19
  prompt, history, max_new_tokens=512, repetition_penalty=1.2, temperature=0.5, top_p=1, top_k=1, seed=42