mmcv/mmcv/utils/parrots_jit.py
Kai Chen c1acedcc3b
Remove pytest from the runtime dependencies (#763)
* update the CI to avoid extra dependencies

* remove the dependency of pytest

* fix typo

* remove skip_no_parrots from __init__
2020-12-28 13:23:48 +08:00

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