Skip to content

GitLab

Commit 8759f303 authored by Ashok M's avatar Ashok M Committed by Soumith Chintala
Browse files

Added test for save image in utils (#847)

parent 6c77fdf6
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 0 deletions
test/test_utils.py
View file @ 8759f303
import os
import torch import torch
import torchvision.utils as utils import torchvision.utils as utils
import unittest import unittest
...@@ -35,6 +36,13 @@ class Tester(unittest.TestCase): ...@@ -35,6 +36,13 @@ class Tester(unittest.TestCase):
assert torch.equal(norm_max, rounded_grid_max), 'Normalized max is not equal to 1' assert torch.equal(norm_max, rounded_grid_max), 'Normalized max is not equal to 1'
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):
t = torch.rand(2, 3, 64, 64)
file_name = 'test_image.png'
utils.save_image(t, file_name)
assert os.path.exists(file_name), 'The image is not present after save'
os.remove(file_name)
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.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