Limited JIT enabling?

jit
1

What I want to do, is run a script with PYTORCH_JIT=0, but still have some free functions optimized (i.e. opt-in mode). My quick test suggests that the following works:

torch.jit._state.enable()
@torch.jit.script
def f(x):
	return x.square()
torch.jit._state.disable() #TODO:restore

Is this reliable or some global objects will be messed up?

2

Apparently, fuser doesn’t work this way. I’ll just resort to skipping some compilations then.