Add digits dataset

braun-steven · Mar 12, 2024 · b4b7fda · b4b7fda
1 parent 593aad8
commit b4b7fda
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/simple_einet/data.py b/simple_einet/data.py
@@ -115,6 +115,7 @@ def get_data_shape(dataset_name: str) -> Shape:
             "flowers": (3, 32, 32),
             "tiny-imagenet": (3, 32, 32),
             "lfw": (3, 32, 32),
+            "digits": (1, 8, 8),
         }[dataset_name]
     )