"vscode:/vscode.git/clone" did not exist on "5a6d52ddae72ced116a36f0b8e28560bc712b54d"
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
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