From 1b9025a4bb40a9063c2643c6abe7fd4703fdde88 Mon Sep 17 00:00:00 2001 From: wxywb Date: Tue, 7 Feb 2023 05:59:07 +0000 Subject: [PATCH] use the corresponding device in the save method. Signed-off-by: wxywb --- clip.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/clip.py b/clip.py index 5ffae49..ec05a33 100644 --- a/clip.py +++ b/clip.py @@ -193,10 +193,10 @@ class Clip(NNOperator): h = sz['height'] w = sz['width'] dummy_input = Image.new('RGB', (w, h), color = 'red') - inputs = self.processor(images=dummy_input, return_tensors='pt') # a dictionary + inputs = self.processor(images=dummy_input, return_tensors='pt').to(self.device) # a dictionary elif self.modality == 'text': dummy_input = 'dummy' - inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt') # a dictionary + inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt').to(self.device) # a dictionary else: raise ValueError("modality[{}] not implemented.".format(self.modality)) @@ -235,6 +235,5 @@ class Clip(NNOperator): opset_version=14, ) else: - pass raise NotImplementedError