GPT2 torch.compile fix (#1434)

huggingface · Oct 18, 2024 · 8b61e8b · 8b61e8b
1 parent f4cb594
commit 8b61e8b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py
@@ -177,7 +177,7 @@ def patch_scoped_linear_all_reduce(model):
 
 
 def get_torch_compiled_model(model):
-    if model.config.model_type in ["gpt_bigcode", "mpt", "bloom"]:
+    if model.config.model_type in ["gpt_bigcode", "mpt", "bloom", "gpt2"]:
         model.transformer = torch.compile(
             model.transformer, backend="hpu_backend", options={"keep_input_mutations": True}
         )