autocast

dnth · Nov 8, 2024 · 3b28a2c · 3b28a2c
1 parent dafbb9d
commit 3b28a2c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/xinfer/transformers/llama32.py b/xinfer/transformers/llama32.py
@@ -126,7 +126,9 @@ def infer(self, image: str, text: str, **generate_kwargs) -> Result:
             self.model.device
         )
 
-        with torch.inference_mode():
+        with torch.inference_mode(), torch.amp.autocast(
+            device_type=self.device, dtype=self.dtype
+        ):
             output = self.model.generate(**inputs, **generate_kwargs)
 
         return Result(text=self.processor.decode(output[0], skip_special_tokens=True))
@@ -143,7 +145,9 @@ def infer_batch(
             self.model.device
         )
 
-        with torch.inference_mode():
+        with torch.inference_mode(), torch.amp.autocast(
+            device_type=self.device, dtype=self.dtype
+        ):
             outputs = self.model.generate(**inputs, **generate_kwargs)
 
         return [