From cad182d29a2f3ad3ed7550b258025f3243981464 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Fri, 28 Feb 2025 18:30:16 -0500
Subject: [PATCH] fix torch compile/dynamo for Gemma2

---
 library/strategy_lumina.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/library/strategy_lumina.py b/library/strategy_lumina.py
index c9e65423..b4c94106 100644
--- a/library/strategy_lumina.py
+++ b/library/strategy_lumina.py
@@ -97,7 +97,8 @@ class LuminaTextEncodingStrategy(TextEncodingStrategy):
                 hidden_states, input_ids, attention_masks
         """
         text_encoder = models[0]
-        assert isinstance(text_encoder, Gemma2Model)
+        # Check model or torch dynamo OptimizedModule
+        assert isinstance(text_encoder, Gemma2Model) or isinstance(text_encoder._orig_mod, Gemma2Model), f"text encoder is not Gemma2Model {text_encoder.__class__.__name__}"
         input_ids, attention_masks = tokens
 
         outputs = text_encoder(