diff --git a/cog.yaml b/cog.yaml
index 7ca0216..4aa3f59 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -6,7 +6,7 @@ build:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
   python_packages:
-    - "min-dalle==0.3.1"
+    - "min-dalle==0.3.3"
   run:
     - pip install torch==1.11.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
 
diff --git a/min_dalle/min_dalle.py b/min_dalle/min_dalle.py
index d8ba51c..686659e 100644
--- a/min_dalle/min_dalle.py
+++ b/min_dalle/min_dalle.py
@@ -170,7 +170,7 @@ class MinDalle:
         self, 
         text: str, 
         seed: int,
-        grid_size: int,
+        image_count: int,
         log2_mid_count: int,
         log2_k: int = 6,
         log2_supercondition_factor: int = 3,
@@ -202,7 +202,7 @@ class MinDalle:
             encoder_state, attention_mask, attention_state, image_tokens = ( 
                 self.decoder.decode_initial(
                     seed, 
-                    grid_size ** 2, 
+                    image_count, 
                     text_tokens, 
                     encoder_state
                 )
@@ -242,7 +242,7 @@ class MinDalle:
         images_stream = self.generate_images_stream(
             text, 
             seed,
-            grid_size,
+            grid_size ** 2,
             log2_mid_count,
             log2_k,
             log2_supercondition_factor,
@@ -256,7 +256,7 @@ class MinDalle:
         self, 
         text: str,
         seed: int = -1,
-        grid_size: int = 1,
+        image_count: int = 1,
         log2_k: int = 6,
         log2_supercondition_factor: int = 3,
         is_verbose: bool = False
@@ -265,7 +265,7 @@ class MinDalle:
         images_stream = self.generate_images_stream(
             text,
             seed,
-            grid_size,
+            image_count,
             log2_mid_count,
             log2_k,
             log2_supercondition_factor,
diff --git a/setup.py b/setup.py
index 35b1b24..ebf9cd9 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ setuptools.setup(
     name='min-dalle',
     description = 'min(DALL·E)',
     # long_description=(Path(__file__).parent / "README.rst").read_text(),
-    version='0.3.1',
+    version='0.3.3',
     author='Brett Kuprel',
     author_email='brkuprel@gmail.com',
     url='https://github.com/kuprel/min-dalle',