SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix device of masks in tests (#27887)

Browse Source 
fix device of mask in tests
This commit is contained in:
fxmarty
2023-12-07 13:34:43 +01:00
committed by GitHub
parent fc71e815f6
commit c99f254763
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/models/llama/test_modeling_llama.py
View File

@@ -104,7 +104,7 @@ class LlamaModelTester:
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:

2
tests/models/mistral/test_modeling_mistral.py
View File

@@ -107,7 +107,7 @@ class MistralModelTester:
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:

2
tests/models/persimmon/test_modeling_persimmon.py
View File

@@ -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: