Spaces:

ssalb
/

story_generator

Paused

ssalb commited on Jan 2

Commit

b990ec8

verified ·

1 Parent(s): dd2160b

Update story_beam_search/beam_search.py

Files changed (1) hide show

story_beam_search/beam_search.py CHANGED Viewed

@@ -93,7 +93,7 @@ class BeamSearchGenerator:
                 top_k=self.config.top_k,
                 top_p=self.config.top_p,
                 do_sample=True,
-            )
         stories = []
         for output in outputs:

                 top_k=self.config.top_k,
                 top_p=self.config.top_p,
                 do_sample=True,
+            ).to(self.device)
         stories = []
         for output in outputs: