Fix device of masks in tests (huggingface#27887)

fix device of mask in tests
NTHU-ML-2023-team19 · Dec 7, 2023 · e332726 · e332726
1 parent 5056ac0
commit e332726
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/tests/models/llama/test_modeling_llama.py b/tests/models/llama/test_modeling_llama.py
@@ -104,7 +104,7 @@ def prepare_config_and_inputs(self):
 
  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:

diff --git a/tests/models/mistral/test_modeling_mistral.py b/tests/models/mistral/test_modeling_mistral.py
@@ -107,7 +107,7 @@ def prepare_config_and_inputs(self):
 
  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:

diff --git a/tests/models/persimmon/test_modeling_persimmon.py b/tests/models/persimmon/test_modeling_persimmon.py
@@ -104,7 +104,7 @@ def prepare_config_and_inputs(self):
 
  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: