Commit 478e456e authored by Aymeric Augustin's avatar Aymeric Augustin
Browse files

Use a random temp dir for writing file in tests.

parent 12726f85
...@@ -18,7 +18,7 @@ from __future__ import print_function ...@@ -18,7 +18,7 @@ from __future__ import print_function
import copy import copy
import sys import sys
import os import os.path
import shutil import shutil
import tempfile import tempfile
import json import json
...@@ -222,16 +222,18 @@ class CommonTestCases: ...@@ -222,16 +222,18 @@ class CommonTestCases:
except RuntimeError: except RuntimeError:
self.fail("Couldn't trace module.") self.fail("Couldn't trace module.")
try: with TemporaryDirectory() as tmp_dir_name:
torch.jit.save(traced_gpt2, "traced_model.pt") pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt")
except RuntimeError:
self.fail("Couldn't save module.")
try: try:
loaded_model = torch.jit.load("traced_model.pt") torch.jit.save(traced_gpt2, pt_file_name)
os.remove("traced_model.pt") except Exception:
except ValueError: self.fail("Couldn't save module.")
self.fail("Couldn't load module.")
try:
loaded_model = torch.jit.load(pt_file_name)
except Exception:
self.fail("Couldn't load module.")
model.to(torch_device) model.to(torch_device)
model.eval() model.eval()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment