From 9219034c5d5866b4f87670428633b1179d245d92 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Tue, 14 May 2024 20:44:04 +0800 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Huazhong Ji Former-commit-id: a435e5a0bdd7268c4f1204f99f289ee0b36fd930 --- src/train.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/train.py b/src/train.py index e2609b66..098ec1b5 100644 --- a/src/train.py +++ b/src/train.py @@ -1,6 +1,4 @@ import os -import torch -import torch_npu from torch_npu.contrib import transfer_to_npu from llmtuner.train.tuner import run_exp @@ -15,6 +13,7 @@ def _mp_fn(index): if __name__ == "__main__": - use_jit_compile = os.getenv('JIT_COMPILE', 'False').lower() in ['true', '1'] - torch.npu.set_compile_mode(jit_compile=use_jit_compile) + if is_torch_npu_available(): + use_jit_compile = os.getenv('JIT_COMPILE', 'False').lower() in ['true', '1'] + torch.npu.set_compile_mode(jit_compile=use_jit_compile) main()