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: