Update chatbot_tutorial.py (pytorch#1535)

Fix lengths to cpu. See issue pytorch#1300 Co-authored-by: Brian Johnson <brianjo@fb.com>
mthrok · May 25, 2021 · fd97ed4 · fd97ed4
1 parent ee37770
commit fd97ed4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/beginner_source/chatbot_tutorial.py b/beginner_source/chatbot_tutorial.py
@@ -1207,7 +1207,7 @@ def evaluate(encoder, decoder, searcher, voc, sentence, max_length=MAX_LENGTH):
     input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)
     # Use appropriate device
     input_batch = input_batch.to(device)
-    lengths = lengths.to(device)
+    lengths = lengths.to("cpu")
     # Decode sentence with searcher
     tokens, scores = searcher(input_batch, lengths, max_length)
     # indexes -> words