catplane

mariogeiger · Nov 13, 2019 · 8a20719 · 8a20719
1 parent f847acb
commit 8a20719
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dataset.py b/dataset.py
@@ -112,6 +112,9 @@ def get_normalized_dataset(dataset, seed):
     elif dataset == "cifar_shipbird":
         tr = [(x, y) for x, y in torchvision.datasets.CIFAR10('~/.torchvision/datasets/CIFAR10', train=True, download=True, transform=transform) if y in [8, 2]]
         te = [(x, y) for x, y in torchvision.datasets.CIFAR10('~/.torchvision/datasets/CIFAR10', train=False, transform=transform) if y in [8, 2]]
+    elif dataset == "cifar_catplane":
+        tr = [(x, y) for x, y in torchvision.datasets.CIFAR10('~/.torchvision/datasets/CIFAR10', train=True, download=True, transform=transform) if y in [3, 0]]
+        te = [(x, y) for x, y in torchvision.datasets.CIFAR10('~/.torchvision/datasets/CIFAR10', train=False, transform=transform) if y in [3, 0]]
     else:
         raise ValueError("unknown dataset")