diff --git a/autogpt/commands/image_gen.py b/autogpt/commands/image_gen.py
index 60cdaec0..962c41fb 100644
--- a/autogpt/commands/image_gen.py
+++ b/autogpt/commands/image_gen.py
@@ -77,12 +77,13 @@ def generate_image_with_hf(prompt: str, filename: str) -> str:
     return f"Saved to disk:{filename}"
 
 
-def generate_image_with_dalle(prompt: str, filename: str) -> str:
+def generate_image_with_dalle(prompt: str, filename: str, size: int) -> str:
     """Generate an image with DALL-E.
 
     Args:
         prompt (str): The prompt to use
         filename (str): The filename to save the image to
+        size (int): The size of the image
 
     Returns:
         str: The filename of the image