Limited JIT enabling?

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:

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

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

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