fix skip condition (#1956)

Lightning-AI · Mar 10, 2025 · 567e256 · 567e256
1 parent 9c59263
commit 567e256
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/litgpt/generate/base.py b/litgpt/generate/base.py
@@ -176,7 +176,7 @@ def generate_fn(
     input_pos = torch.arange(0, prompt_size, device=device, dtype=torch.int64)
     # input_pos_maxp1 introduces data-dependent shapes and control flow.
     # We want to skip if ThunderModules are involved, either directly or wrapped in LightningModule etc.
-    if any(m.__class__.__name__ != 'ThunderModule' for m in model.modules()):
+    if not any(m.__class__.__name__ == 'ThunderModule' for m in model.modules()):
         input_pos_maxp1 = torch.tensor(prompt_size, device=device)
     else:
         input_pos_maxp1 = None