Spaces:
Sleeping
Sleeping
File size: 1,031 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
"""
APIs related to torch.compile which lazily import torch._dynamo to avoid
circular dependencies.
"""
import functools
def _disable_dynamo(fn=None, recursive=True):
"""
This API should be only used inside torch, external users should still use
torch._dynamo.disable. The main goal of this API is to avoid circular
imports issues that is common while using _dynamo.disable inside torch
itself.
This API avoids it by lazily importing torch._dynamo from the import time to
the invocation of the decorated function.
"""
if fn is not None:
@functools.wraps(fn)
def inner(*args, **kwargs):
import torch._dynamo
return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
return inner
else:
# decorator usage like @_disable_dynamo(recursive=False). The resulting
# object expects the original decorated function as the arg.
return functools.partial(_disable_dynamo, recursive=recursive)
|