From eaaf1c1766cc0befed1bb8a24ffaf3847796f148 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 10 Sep 2025 20:25:46 +0200 Subject: [PATCH] additionally --- .../models/gemma/modeling_gemma.py | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/modeling_gemma.py b/src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/modeling_gemma.py index b596bcad5..a3f6b5325 100644 --- a/src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/modeling_gemma.py +++ b/src/lerobot/policies/pi0_openpi/transformers_replace/models/gemma/modeling_gemma.py @@ -46,13 +46,22 @@ logger = logging.get_logger(__name__) # Workaround for Python 3.10+ UnionType compatibility with transformers auto_docstring -def safe_auto_docstring(func): +def safe_auto_docstring(func=None, **kwargs): """Auto docstring decorator that handles Python 3.10+ UnionType gracefully.""" - try: - return auto_docstring(func) - except (AttributeError, TypeError): - # If auto_docstring fails due to UnionType, just return the function unchanged - return func + + def decorator(f): + try: + return auto_docstring(f, **kwargs) if kwargs else auto_docstring(f) + except (AttributeError, TypeError): + # If auto_docstring fails due to UnionType, just return the function unchanged + return f + + if func is None: + # Called with arguments, return the decorator + return decorator + else: + # Called without arguments, apply directly + return decorator(func) class GemmaRMSNorm(nn.Module):