Add sample batch size for Lumina

sdbds · Feb 24, 2025 · 48e7da2 · 48e7da2
1 parent ba725a8
commit 48e7da2
Show file tree

Hide file tree

Showing 3 changed files with 199 additions and 104 deletions.
diff --git a/library/lumina_models.py b/library/lumina_models.py
@@ -880,8 +880,8 @@ def __init__(
         self.n_heads = n_heads
 
         self.gradient_checkpointing = False
-        self.cpu_offload_checkpointing = False
-        self.blocks_to_swap = None
+        self.cpu_offload_checkpointing = False # TODO: not yet supported
+        self.blocks_to_swap = None # TODO: not yet supported
 
     @property
     def device(self):
@@ -982,8 +982,8 @@ def patchify_and_embed(
 
         l_effective_cap_len = cap_mask.sum(dim=1).tolist()
         encoder_seq_len = cap_mask.shape[1]
-
         image_seq_len = (height // self.patch_size) * (width // self.patch_size)
+
         seq_lengths = [cap_seq_len + image_seq_len for cap_seq_len in l_effective_cap_len]
         max_seq_len = max(seq_lengths)