o ÓÙ¾išã@s(dZddlZddlZdd„Zdd„ZdS)zP JIT scripting/tracing utils Hacked together by / Copyright 2020 Ross Wightman éNcCs<ttjdƒs Jdƒ‚tj d¡tj d¡tj d¡dS)zÝ Set JIT executor to legacy w/ support for op fusion This is hopefully a temporary need in 1.5/1.5.1/1.6 to restore performance due to changes in the JIT executor. These API are not supported so could change. Ú_jit_set_profiling_executorzOld JIT behavior doesn't exist!FTN)ÚhasattrÚtorchÚ_CrÚ_jit_set_profiling_modeÚ_jit_override_can_fuse_on_gpu©rrúB/home/ubuntu/.local/lib/python3.10/site-packages/timm/utils/jit.pyÚset_jit_legacy s  r cCsj|dkr6tj d¡tj d¡tj d¡tj d¡tj d¡z tj d¡WdSty5YdSw|dks>|dkrjtj d¡tj d¡tj d¡tj d¡z tj d¡WdStyiYdSw|dksr|dkr­dt j d <tj d¡tj d¡tj d¡tj  ¡tj  ¡tj d¡tj d¡tj  d¡tj d¡dSJd |›d ƒ‚) NÚteTFÚoldÚlegacyÚnvfuserÚnvfÚ1Ú PYTORCH_NVFUSER_DISABLE_FALLBACKzInvalid jit fuser (ú))rrrrÚ_jit_override_can_fuse_on_cpurÚ_jit_set_texpr_fuser_enabledÚ_jit_set_nvfuser_enabledÚ ExceptionÚosÚenvironÚ_jit_can_fuse_on_cpuÚ_jit_can_fuse_on_gpuÚ_jit_set_nvfuser_guard_mode)Úfuserrrr Ú set_jit_fusersB      ÿ     ÿ         r)Ú__doc__rrr rrrrr Ús