Enable JIT only when cuda is available

#84
by statelesshz - opened
Files changed (1) hide show
  1. modeling_chatglm.py +1 -1
modeling_chatglm.py CHANGED
@@ -29,7 +29,7 @@ from .configuration_chatglm import ChatGLMConfig
29
 
30
  # flags required to enable jit fusion kernels
31
 
32
- if sys.platform != 'darwin':
33
  torch._C._jit_set_profiling_mode(False)
34
  torch._C._jit_set_profiling_executor(False)
35
  torch._C._jit_override_can_fuse_on_cpu(True)
 
29
 
30
  # flags required to enable jit fusion kernels
31
 
32
+ if sys.platform != 'darwin' and torch.cuda.is_available():
33
  torch._C._jit_set_profiling_mode(False)
34
  torch._C._jit_set_profiling_executor(False)
35
  torch._C._jit_override_can_fuse_on_cpu(True)