diff --git a/gemma/gm/nn/gemma4/_moe.py b/gemma/gm/nn/gemma4/_moe.py
index 25fc0d0d..266dcc4c 100644
--- a/gemma/gm/nn/gemma4/_moe.py
+++ b/gemma/gm/nn/gemma4/_moe.py
@@ -389,7 +389,7 @@ def __call__(self, x, unnormalized_x=None):
         the router applies only its own router_norm (avoiding a double norm).
 
     Returns:
-      Output of the MoE module.
+      Output of the MoERagged module.
     """
     if unnormalized_x is None:
       unnormalized_x = x