Add missing DEFAULT_TRAIN_BSIZE field

Author: ostrowskimarcin

torchbenchmark/models/cm3leon_generate/__init__.py

@@ -8,6 +8,7 @@
 class Model(BenchmarkModel):
     task = NLP.LANGUAGE_MODELING
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1 
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):
         super().__init__(

torchbenchmark/models/doctr_det_predictor/__init__.py

@@ -13,6 +13,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.DETECTION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

torchbenchmark/models/doctr_reco_predictor/__init__.py

@@ -13,6 +13,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.DETECTION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

torchbenchmark/models/llama/__init__.py

@@ -12,6 +12,7 @@
 class Model(BenchmarkModel):
     task = NLP.LANGUAGE_MODELING
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):
         super().__init__(

torchbenchmark/models/pyhpc_equation_of_state/__init__.py

@@ -40,6 +40,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data point: i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

torchbenchmark/models/pyhpc_isoneutral_mixing/__init__.py

@@ -131,6 +131,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data-point when i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 
     def __init__(self, test, device, batch_size=None, extra_args=[]):

torchbenchmark/models/pyhpc_turbulent_kinetic_energy/__init__.py

@@ -129,6 +129,7 @@ class Model(BenchmarkModel):
     # Source: https://github.com/dionhaefner/pyhpc-benchmarks/blob/650ecc650e394df829944ffcf09e9d646ec69691/run.py#L25
     # Pick data-point when i = 20, size = 1048576
     DEFAULT_EVAL_BSIZE = 1048576
+    DEFAULT_TRAIN_BSIZE = 1048576
     ALLOW_CUSTOMIZE_BSIZE = False
     CANNOT_SET_CUSTOM_OPTIMIZER = True
 

torchbenchmark/models/sam/__init__.py

@@ -17,6 +17,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.SEGMENTATION
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=1, extra_args=[]):
         super().__init__(

torchbenchmark/models/sam_fast/__init__.py

@@ -18,6 +18,7 @@
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.SEGMENTATION
     DEFAULT_EVAL_BSIZE = 32
+    DEFAULT_TRAIN_BSIZE = 32
 
     def __init__(self, test, device, batch_size=1, extra_args=[]):
         super().__init__(

torchbenchmark/models/simple_gpt/__init__.py

@@ -11,6 +11,7 @@
 class Model(BenchmarkModel):
     task = NLP.GENERATION
     DEFAULT_EVAL_BSIZE = 1
+    DEFAULT_TRAIN_BSIZE = 1
 
     def validate_environment(self):
         if not torch.cuda.is_available() or "cuda" not in self.device: