diff --git a/sdk/python/kfp/dsl/_container_op.py b/sdk/python/kfp/dsl/_container_op.py
index 99d39b8def7..2d8299b36cf 100644
--- a/sdk/python/kfp/dsl/_container_op.py
+++ b/sdk/python/kfp/dsl/_container_op.py
@@ -769,9 +769,9 @@ def apply(self, mod_func):
           from kfp.gcp import use_gcp_secret
           task = (
             train_op(...)
-              .set_memory_request('1GB')
+              .set_memory_request('1G')
               .apply(use_gcp_secret('user-gcp-sa'))
-              .set_memory_limit('2GB')
+              .set_memory_limit('2G')
           )
         """
         return mod_func(self) or self