optimised

sets/CelebA.py

@@ -0,0 +1,20 @@
+from torchvision import datasets
+from torch.utils.data import Dataset
+import torch
+from .Data import Data
+
+class CelebA(Data):
+    def __init__(self, resolution: int = 224, sample_size = 30):
+        super().__init__(resolution, sample_size = sample_size)
+
+    def get_set(self):
+        set = datasets.CelebA(
+            root = "../data", 
+            split='all', 
+            target_type='identity', 
+            download=False, 
+            transform=None
+        )
+        # set the target first
+        self.target = set.identity
+        return set