UX code change (#1764)

Co-authored-by: Libin Tang <[email protected]> Co-authored-by: regisss <[email protected]>
huggingface · Feb 12, 2025 · 35fa12d · 35fa12d
1 parent 97beef8
commit 35fa12d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py
@@ -315,6 +315,9 @@ def setup_model(args, model_dtype, model_kwargs, logger):
 
     if args.torch_compile:
         model = get_torch_compiled_model(model, logger)
+        assert "PT_HPU_LAZY_MODE" in os.environ and os.environ["PT_HPU_LAZY_MODE"] == "0", (
+            "Please set PT_HPU_LAZY_MODE=0 on command line when using `--torch_compile`"
+        )
         # if args.assistant_model is not None:
         #     assistant_model = get_torch_compiled_model(assistant_model, logger)
     return model, assistant_model