make style fixes

dvarshney-habana · Aug 2, 2024 · 00c84c1 · 00c84c1
1 parent 191289e
commit 00c84c1
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py
@@ -177,10 +177,12 @@ def patch_scoped_linear_all_reduce(model):
 
 
 def get_torch_compiled_model(model):
-    if  model.config.model_type in ['gpt_bigcode']:
-      model.transformer = torch.compile(model.transformer, backend="hpu_backend", options={"keep_input_mutations": True})
+    if model.config.model_type in ["gpt_bigcode"]:
+        model.transformer = torch.compile(
+            model.transformer, backend="hpu_backend", options={"keep_input_mutations": True}
+        )
     else:
-      model.model = torch.compile(model.model, backend="hpu_backend", options={"keep_input_mutations": True})
+        model.model = torch.compile(model.model, backend="hpu_backend", options={"keep_input_mutations": True})
     return model