luna-playground

Sleeping

terryyz commited on Mar 12

Commit

d46f31c

•

1 Parent(s): e90f65d

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -180,7 +180,7 @@ def generate(
         model = model_map["CS Evo Prefix"]
     else:
         output = ""
-    # model.to(device)
     input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
     # generated_ids = model.generate(**input_ids
     generated_ids = model.generate(**input_ids)#, **generate_kwargs)

         model = model_map["CS Evo Prefix"]
     else:
         output = ""
+    model.to(device)
     input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
     # generated_ids = model.generate(**input_ids
     generated_ids = model.generate(**input_ids)#, **generate_kwargs)