Update story_beam_search/beam_search.py

story_beam_search/beam_search.py

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