Add missing DEFAULT_TRAIN_BSIZE field

pytorch · Jan 13, 2025 · f421586 · f421586
1 parent 9028c93
commit f421586
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchbenchmark/models/cm3leon_generate/__init__.py b/torchbenchmark/models/cm3leon_generate/__init__.py
@@ -7,6 +7,7 @@
 
 class Model(BenchmarkModel):
     task = NLP.LANGUAGE_MODELING
+    DEFAULT_TRAIN_BSIZE = 1
     DEFAULT_EVAL_BSIZE = 1
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):