lower depth

kyegomez · Oct 2, 2023 · 1268373 · 1268373
1 parent d34c441
commit 1268373
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -32,8 +32,8 @@
     patch_size=14,
     num_classes=10,  # CIFAR-10 has 10 classes
     dim=6144,
-    depth=48,
-    heads=48,
+    depth=12,
+    heads=12,
     mlp_dim=2048,
     dropout=0.1,
     emb_dropout=0.1