From b4b7fdaf92eb8b2e9c046ab18ff5d65fa3dbd003 Mon Sep 17 00:00:00 2001
From: Steven Braun <steven.braun.mz@gmail.com>
Date: Tue, 12 Mar 2024 07:23:56 +0100
Subject: [PATCH] Add digits dataset

---
 simple_einet/data.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/simple_einet/data.py b/simple_einet/data.py
index 8cfc60c..de2ec01 100644
--- 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]
     )