ydshieh
/
vit-gpt2-coco-en
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update pipeline.py

This commit is contained in:
Yih-Dar SHIEH 2021-10-25 08:13:14 +00:00 committed by huggingface-web
parent 25a7779507
commit 117d46a115
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pipeline.py
View File

@ -20,7 +20,7 @@ class PreTrainedPipeline():
max_length = 16
num_beams = 4
# self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, return_dict_in_generate=True}
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True}
self.model.to("cpu")
self.model.eval()