diff --git a/torchbenchmark/models/hf_Falcon_7b/__init__.py b/torchbenchmark/models/hf_Falcon_7b/__init__.py index 53a163c49c..2e18da823c 100644 --- a/torchbenchmark/models/hf_Falcon_7b/__init__.py +++ b/torchbenchmark/models/hf_Falcon_7b/__init__.py @@ -4,7 +4,8 @@ class Model(HuggingFaceModel): task = NLP.LANGUAGE_MODELING # Published training batch size is 2304: see https://huggingface.co/tiiuae/falcon-7b/blob/main/README.md - DEFAULT_TRAIN_BSIZE = 2304 + # Setting to default value of 4 for feasibility + DEFAULT_TRAIN_BSIZE = 4 DEFAULT_EVAL_BSIZE = 1 def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):