allow patch_batch_size to be adjusted in the forward() method

facebookresearch · Jan 18, 2025 · cff0dcb · cff0dcb
1 parent 175fce6
commit cff0dcb
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/bytelatent/data/patcher.py b/bytelatent/data/patcher.py
@@ -490,6 +490,7 @@ def patch(
         preds: torch.Tensor | None = None,
         entropies: torch.Tensor | None = None,
         threshold: float = None,
+        patching_batch_size: int | None = None,
     ) -> torch.Tensor:
         """
         tokens: 2D tensor of shape [batch_size, seq_len] that needs to be patched
@@ -539,7 +540,11 @@ def patch(
                 scores = calculate_entropies(
                     tokens,
                     self.entropy_model,
-                    self.patching_batch_size,
+                    (
+                        patching_batch_size
+                        if patching_batch_size is not None
+                        else self.patching_batch_size
+                    ),
                     self.device,
                 )
             if self.log_time: