fix device handling

custom_st.py

@@ -27,7 +27,6 @@ class Transformer(nn.Module):
         tokenizer_args: Optional[Dict[str, Any]] = None,
         config_args: Optional[Dict[str, Any]] = None,
         cache_dir: Optional[str] = None,
-        device: str = 'cpu',
         backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
         **kwargs,
     ) -> None:
@@ -38,7 +37,6 @@ class Transformer(nn.Module):
                 f'Backend \'{backend}\' is not supported, please use \'torch\' instead'
             )
         
-        self.device = device
         self.dimension = dimension
         self.max_pixels = max_pixels
         self.min_pixels = min_pixels
@@ -160,15 +158,13 @@ class Transformer(nn.Module):
     def tokenize(self, texts: List[Union[str, Image.Image]], padding: str = 'longest') -> Dict[str, torch.Tensor]:
         processed_texts, processed_images = self._process_input(texts)
         
-        inputs = self.processor(
+        return self.processor(
             text=processed_texts,
             images=processed_images,
             videos=None,
             padding=padding,
             return_tensors='pt'
         )
-        
-        return {k: v.to(self.device) for k, v in inputs.items()}
 
     def save(self, output_path: str, safe_serialization: bool = True) -> None:
         """Save the model, tokenizer and processor to the given path."""