logo
Browse Source

Debug decorator

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 2 years ago
parent
commit
21b72075da
  1. 2
      test_onnx.py
  2. 2
      timm_image.py

2
test_onnx.py

@ -57,6 +57,8 @@ for name in models:
status = [name] + ['fail'] * 5
try:
out0 = op(data.cpu().detach().numpy().squeeze(0))
assert out0.shape
out1 = op.model(data).detach().numpy()
logger.info('OP LOADED.')
status[1] = 'success'

2
timm_image.py

@ -44,7 +44,7 @@ log = logging.getLogger('timm_op')
def torch_no_grad(f):
def wrap(*args, **kwargs):
with torch.no_grad():
f(*args, **kwargs)
return f(*args, **kwargs)
return wrap

Loading…
Cancel
Save