test torch JIT

app.py

@@ -12,6 +12,12 @@ torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dino_v2_model = AutoModel.from_pretrained("./dinov2-large").to(torch_device)
 dino_v2_image_processor = AutoImageProcessor.from_pretrained("./dinov2-large")
 
+# Provide a sample input for tracing
+sample_input = dino_v2_image_processor(
+    images=Image.new("RGB", (64, 64)), return_tensors="pt"
+).to(torch_device)
+traced_dino_v2_model = torch.jit.trace(dino_v2_model, sample_input["pixel_values"])
+
 
 def process_image(image):
     """
@@ -46,7 +52,9 @@ def process_image(image):
         inputs = dino_v2_image_processor(images=image, return_tensors="pt").to(
             torch_device
         )
-        outputs = dino_v2_model(**inputs)
+
+        # Use the traced model for inference
+        outputs = traced_dino_v2_model(**inputs)
 
     # Normalize the features before search, whatever that means
     embeddings = outputs.last_hidden_state