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

Fix torchscript tests (#13701)

Browse Source
This commit is contained in:
Lysandre Debut
2021-09-22 19:02:54 -04:00
committed by GitHub
parent 5b57075449
commit ca257a06cc
3 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_modeling_distilbert.py
View File

@@ -273,7 +273,7 @@ class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):
with tempfile.TemporaryDirectory() as tmp:
torch.jit.save(traced_model, os.path.join(tmp, "traced_model.pt"))
loaded = torch.jit.load(os.path.join(tmp, "bert.pt"), map_location=torch_device)
loaded = torch.jit.load(os.path.join(tmp, "traced_model.pt"), map_location=torch_device)
loaded(inputs_dict["input_ids"].to(torch_device), inputs_dict["attention_mask"].to(torch_device))