lower params

kyegomez · Oct 2, 2023 · 71213c8 · 71213c8
1 parent 1268373
commit 71213c8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/train.py b/train.py
@@ -31,9 +31,9 @@
     image_size=224,
     patch_size=14,
     num_classes=10,  # CIFAR-10 has 10 classes
-    dim=6144,
-    depth=12,
-    heads=12,
+    dim=1024,
+    depth=6,
+    heads=16,
     mlp_dim=2048,
     dropout=0.1,
     emb_dropout=0.1