burtenshaw HF Staff commited on
Commit
602b318
·
verified ·
1 Parent(s): 3cf0ef6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -28,7 +28,8 @@ def generate(prompt, history):
28
  conversation,
29
  add_generation_prompt=True,
30
  tokenize=True,
31
- return_tensors="pt"
 
32
  ).to(device)
33
 
34
  with torch.no_grad():
 
28
  conversation,
29
  add_generation_prompt=True,
30
  tokenize=True,
31
+ return_tensors="pt",
32
+ return_dict=True,
33
  ).to(device)
34
 
35
  with torch.no_grad():