BatchEncoding.to with device with tests (#9584)

2021-01-14 07:57:58 -05:00
parent 8bf27075a2
commit 280db79ac1
2 changed files with 11 additions and 1 deletions
--- a/tests/test_tokenization_common.py
+++ b/tests/test_tokenization_common.py
@@ -1704,6 +1704,10 @@ class TokenizerTesterMixin:
                first_ten_tokens = list(tokenizer.get_vocab().keys())[:10]
                sequence = " ".join(first_ten_tokens)
                encoded_sequence = tokenizer.encode_plus(sequence, return_tensors="pt")
+
+                # Ensure that the BatchEncoding.to() method works.
+                encoded_sequence.to(model.device)
+
                batch_encoded_sequence = tokenizer.batch_encode_plus([sequence, sequence], return_tensors="pt")
                # This should not fail