import torch def set_fuser(fuser_name, executor_name): assert fuser_name in ['te', 'old', 'none', 'default'] if fuser_name == 'te': torch._C._jit_set_profiling_executor(True) torch._C._jit_set_profiling_mode(True) torch._C._jit_override_can_fuse_on_cpu(False) torch._C._jit_override_can_fuse_on_gpu(True) torch._C._jit_set_texpr_fuser_enabled(True) elif fuser_name == 'old': torch._C._jit_set_profiling_executor(False) torch._C._jit_set_profiling_mode(False) torch._C._jit_override_can_fuse_on_gpu(True) torch._C._jit_set_texpr_fuser_enabled(False) elif fuser_name == 'none': torch._C._jit_set_profiling_executor(False) torch._C._jit_set_profiling_mode(False) torch._C._jit_override_can_fuse_on_gpu(False) torch._C._jit_override_can_fuse_on_cpu(False) torch._C._jit_set_texpr_fuser_enabled(False) elif fuser_name == 'default': pass # --executor overrides settings of --fuser if executor_name == 'profiling': torch._C._jit_set_profiling_executor(True) torch._C._jit_set_profiling_mode(True) elif executor_name == 'simple': torch._C._jit_set_profiling_executor(True) torch._C._jit_set_profiling_mode(False) elif executor_name == 'legacy': torch._C._jit_set_profiling_executor(False) torch._C._jit_set_profiling_mode(False) elif executor_name == 'default': pass