windows doesn't like opening tempfile twice

luigibonati · Apr 17, 2023 · 54ef4b9 · 54ef4b9
1 parent 33fa009
commit 54ef4b9
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/mlcvs/tests/test_cvs_multitask_multitask.py b/mlcvs/tests/test_cvs_multitask_multitask.py
@@ -14,6 +14,7 @@
 # GLOBAL IMPORTS
 # =============================================================================
 
+import os
 import tempfile
 
 import pytest
@@ -283,8 +284,13 @@ def test_multitask_training(main_cv_name, weights, auxiliary_loss_names, loss_co
     assert x_hat.shape == (x.shape[0], N_CVS)
 
     # Do round-trip through torchscript.
-    with tempfile.NamedTemporaryFile('wb', suffix='.ptc') as f:
-        multi_cv.to_torchscript(file_path=f.name, method='trace')
-        multi_cv_loaded = torch.jit.load(f.name)
+    # This try-finally clause is a workaround for windows not allowing opening temp files twice.
+    try:
+        tmp_file = tempfile.NamedTemporaryFile('wb', suffix='.ptc', delete=False)
+        tmp_file.close()
+        multi_cv.to_torchscript(file_path=tmp_file.name, method='trace')
+        multi_cv_loaded = torch.jit.load(tmp_file.name)
+    finally:
+        os.unlink(tmp_file.name)
     x_hat2 = multi_cv_loaded(x)
     assert torch.allclose(x_hat, x_hat2)
diff --git a/mlcvs/tests/test_cvs_unsupervised_vae.py b/mlcvs/tests/test_cvs_unsupervised_vae.py
@@ -14,6 +14,7 @@
 # GLOBAL IMPORTS
 # =============================================================================
 
+import os
 import tempfile
 
 import pytest
@@ -63,8 +64,13 @@ def test_vae_cv_training(weights):
     assert x_hat.shape == (batch_size, n_cvs)
 
     # Test export to torchscript.
-    with tempfile.NamedTemporaryFile('wb', suffix='.ptc') as f:
-        model.to_torchscript(file_path=f.name, method='trace')
-        model_loaded = torch.jit.load(f.name)
+    # This try-finally clause is a workaround for windows not allowing opening temp files twice.
+    try:
+        tmp_file = tempfile.NamedTemporaryFile('wb', suffix='.ptc', delete=False)
+        tmp_file.close()
+        model.to_torchscript(file_path=tmp_file.name, method='trace')
+        model_loaded = torch.jit.load(tmp_file.name)
+    finally:
+        os.unlink(tmp_file.name)
     x_hat2 = model_loaded(x)
     assert torch.allclose(x_hat, x_hat2)