torch.compiler.disable doesn't disable nested functions (also doesn't work as a context manager) #123771
Description
🐛 Describe the bug
import torch
torch.set_default_device('cuda')
@torch.compile(fullgraph=True)
def f(x):
x = x.cos().cos()
print(x)
return x.cos()
def g(x):
return f(x)
torch.compiler.disable(g)(torch.randn(5))
This still compiles (and throws the error).
Also, the docstring also claims it can be used as a context manager, but this error is thrown when you try to use it that way.
RuntimeError: torch._dynamo.optimize(...) is used with a context manager. Please refer to https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html to use torch._dynamo.optimize(...) as an annotation/decorator.
cc: @anijain2305
Versions
N/A
cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @msaroufim @bdhirsh @anijain2305