diff --git a/clip_vision.py b/clip_vision.py index a7c3e03..273593e 100644 --- a/clip_vision.py +++ b/clip_vision.py @@ -17,9 +17,11 @@ import numpy as np from torchvision import transforms as T from towhee.operator import NNOperator, SharedType +from towhee.dc2 import accelerate from towhee.models import clip +@accelerate class Model: def __init__(self, model_name, device='cpu'): self.model = clip.create_model(model_name=model_name, pretrained=True, device=device).visual