Update model_factory.py

pytorch · Jul 22, 2023 · abc1ed3 · abc1ed3
1 parent 9fbceda
commit abc1ed3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/util/framework/huggingface/model_factory.py b/torchbenchmark/util/framework/huggingface/model_factory.py
@@ -28,7 +28,7 @@
     # see https://huggingface.co/bert-large-cased
     'hf_Bert_large': (512, 512, 'BertConfig(hidden_size=1024, num_hidden_layers=24, num_attention_heads=16)', 'AutoModelForMaskedLM'),
     # default num_hidden_layers=32 but that OOMs, feel free to change this config to something more real
-    'llama_v2' : (512,512, 'LlamaConfig(num_hidden_layers=16)', 'AutoModelForCausalLM'),
+    'llama_v2' : (512,512, 'LlamaConfig()', 'AutoModelForCausalLM'),
 }
 
 cpu_input_slice = {