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?