Skip to content

GitLab

Unverified Commit 174e135d authored by Francisco Massa's avatar Francisco Massa Committed by GitHub
Browse files

Make test_save_image more robust (#1037)

parent 3b8a1403
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 5 deletions
test/test_utils.py
View file @ 174e135d
import os import os
import tempfile
import torch import torch
import torchvision.utils as utils import torchvision.utils as utils
import unittest import unittest
...@@ -37,11 +38,10 @@ class Tester(unittest.TestCase): ...@@ -37,11 +38,10 @@ class Tester(unittest.TestCase):
assert torch.equal(norm_min, rounded_grid_min), 'Normalized min is not equal to 0' assert torch.equal(norm_min, rounded_grid_min), 'Normalized min is not equal to 0'
def test_save_image(self): def test_save_image(self):
with tempfile.NamedTemporaryFile(suffix='.png') as f:
t = torch.rand(2, 3, 64, 64) t = torch.rand(2, 3, 64, 64)
file_name = 'test_image.png' utils.save_image(t, f.name)
utils.save_image(t, file_name) assert os.path.exists(f.name), 'The image is not present after save'
assert os.path.exists(file_name), 'The image is not present after save'
os.remove(file_name)
if __name__ == '__main__': if __name__ == '__main__':
......
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