From 332f44fa43e46df112ca62035905f3a3bfaf77d8 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: 0ac6e73f9971a9310026ddc609b5266cb1639b64 --- 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()