Fix device of masks in tests (#27887)

fix device of mask in tests
2023-12-07 13:34:43 +01:00
parent fc71e815f6
commit c99f254763
3 changed files with 3 additions and 3 deletions
--- a/tests/models/persimmon/test_modeling_persimmon.py
+++ b/tests/models/persimmon/test_modeling_persimmon.py
@@ -104,7 +104,7 @@ class PersimmonModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = torch.tril(torch.ones(self.batch_size, self.seq_length))
+            input_mask = torch.tril(torch.ones(self.batch_size, self.seq_length)).to(torch_device)

        token_type_ids = None
        if self.use_token_type_ids: