test

example.py

@@ -78,7 +78,7 @@ class VideoHighlightDetector:
             tokenize=True,
             return_dict=True,
             return_tensors="pt"
-        ).to(self.device)
+        ).to(self.device, dtype=torch.bfloat16)
         
         outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
         return self.processor.decode(outputs[0], skip_special_tokens=True).lower().split("assistant: ")[1]

src/video_processor/processor.py

@@ -67,7 +67,7 @@ class VideoAnalyzer:
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt"
-            ).to(DEVICE)
+            ).to(DEVICE, dtype=torch.bfloat16)
             
             logger.info(f"Generating IDs")
             # Generate description with increased token limit