Update pipeline.py

pipeline.py

@@ -18,6 +18,12 @@ class PreTrainedPipeline():
         max_length = 16
         num_beams = 4
         self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+        
+        @jax.jit
+        def _generate(pixel_values):
+
+            output_ids = self.model.generate(pixel_values, **self.gen_kwargs).sequences
+            return output_ids
 
         # compile the model
         image_path = os.path.join(path, 'val_000000039769.jpg')
@@ -25,11 +31,7 @@ class PreTrainedPipeline():
         self(image)
         image.close()
 
-    @jax.jit
-    def generate(self, pixel_values):
-
-        output_ids = self.model.generate(pixel_values, **self.gen_kwargs).sequences
-        return output_ids
+        self.generate = _generate
 
     def __call__(self, inputs: "Image.Image") -> List[str]:
         """