diff --git a/awswrangler/_distributed.py b/awswrangler/_distributed.py
index aa0735d38..dffb85816 100644
--- a/awswrangler/_distributed.py
+++ b/awswrangler/_distributed.py
@@ -177,7 +177,7 @@ def get(cls) -> MemoryFormatEnum:
             return cls._enum if cls._enum else cls.get_installed()
 
     @classmethod
-    def set(cls, name: EngineLiteral) -> None:
+    def set(cls, name: MemoryFormatLiteral) -> None:
         """Set the memory format."""
         with cls._lock:
             cls._enum = MemoryFormatEnum[name.upper()]