default BS=9 in handcode_opt bert (tinygrad#6783)

using 54 for 6 gpus now, and 2 is not a good default
chenyuxyz · Sep 27, 2024 · 2fc2689 · 2fc2689
1 parent 9a3f6f3
commit 2fc2689
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/handcode_opt.py b/examples/handcode_opt.py
@@ -33,7 +33,7 @@ def get_sched_bert():
   optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
 
   # fake data
-  BS = getenv("BS", 2)
+  BS = getenv("BS", 9)
   input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
   segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
   attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)