From f0270af71e9270f0c722ea5c29473641e76edd48 Mon Sep 17 00:00:00 2001 From: Ashok M Date: Wed, 10 Apr 2019 20:39:49 +0200 Subject: [PATCH] Added test for save image in utils --- test/test_utils.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test/test_utils.py b/test/test_utils.py index 9a5349fbdb0..07b4b8fd4a1 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -1,3 +1,4 @@ +import os import torch import torchvision.utils as utils import unittest @@ -35,6 +36,13 @@ def test_normalize_in_make_grid(self): 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' + 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__': unittest.main()