Added test_load_data

Author: Seilmast

tests/test_wrappers.py

@@ -28,4 +28,34 @@ def test_load_model():
 
 
 def test_load_data():
-    pass
+    from tempfile import TemporaryDirectory
+
+    import torch as th
+    from torchvision import transforms
+
+    dataset_names = [
+        "usps_0-6",
+        "mnist_0-3",
+        "usps_7-9",
+        "svhn",
+        # 'mnist_4-9' #Uncomment when implemented
+    ]
+
+    trans = transforms.Compose(
+        [
+            transforms.Resize((16, 16)),
+            transforms.ToTensor(),
+        ]
+    )
+
+    with TemporaryDirectory() as tmppath:
+        for name in dataset_names:
+            dataset = load_data(
+                name, train=False, data_path=tmppath, download=True, transform=trans
+            )
+
+            im, lab = dataset.__getitem__(0)
+
+            assert dataset.__len__() != 0
+            assert type(im) == th.Tensor and len(im.size()) == 3
+            assert lab - lab == 0.0