From 505b0b220d830c042175bbf20cc0da643326e44e Mon Sep 17 00:00:00 2001 From: "junjie.jiang" Date: Mon, 27 Mar 2023 17:13:26 +0800 Subject: [PATCH] enable accelerate Signed-off-by: junjie.jiang --- timm_image.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/timm_image.py b/timm_image.py index a5690a3..515f5c2 100644 --- a/timm_image.py +++ b/timm_image.py @@ -23,7 +23,11 @@ from towhee.operator.base import NNOperator, OperatorFlag from towhee.types.arg import arg, to_image_color from towhee import register from towhee.types import Image -# from towhee.dc2 import accelerate +try: + from towhee import accelerate +except: + def accelerate(func): + return func import torch from torch import nn @@ -48,7 +52,7 @@ def torch_no_grad(f): return wrap -# @accelerate +@accelerate class Model: def __init__(self, model_name, device, num_classes): self.device = device