File size: 644 Bytes
01badea
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
d0142c9dcad8f456caa482407960dc9cb5f0e710
diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py
index 280e881..a2f6b5c 100644
--- a/modelopt/torch/utils/dataset_utils.py
+++ b/modelopt/torch/utils/dataset_utils.py
@@ -945,6 +945,8 @@ _NESTED_USE_CACHE_CONFIG_ATTRS = ("text_config",)
 
 def _iter_use_cache_configs(model: torch.nn.Module) -> Iterator[Any]:
     """Yield the top-level config and Step3.7-style nested text config."""
+    if "diffusiongemma" in model.__class__.__name__.lower():
+        return
     seen: set[int] = set()
     config = getattr(model, "config", None)
     if config is None: