mirror of https://github.com/open-mmlab/mmcv.git
37 lines
755 B
Python
37 lines
755 B
Python
from .parrots_wrapper import TORCH_VERSION
|
|
|
|
if TORCH_VERSION == 'parrots':
|
|
from parrots.jit import pat as jit
|
|
else:
|
|
|
|
def jit(func=None,
|
|
check_input=None,
|
|
full_shape=True,
|
|
derivate=False,
|
|
coderize=False,
|
|
optimize=False):
|
|
|
|
def wrapper(func):
|
|
|
|
def wrapper_inner(*args, **kargs):
|
|
return func(*args, **kargs)
|
|
|
|
return wrapper_inner
|
|
|
|
if func is None:
|
|
return wrapper
|
|
else:
|
|
return func
|
|
|
|
|
|
if TORCH_VERSION == 'parrots':
|
|
from parrots.utils.tester import skip_no_elena
|
|
else:
|
|
|
|
def skip_no_elena(func):
|
|
|
|
def wrapper(*args, **kargs):
|
|
return func(*args, **kargs)
|
|
|
|
return wrapper
|