diff --git a/coreml_suite/models.py b/coreml_suite/models.py
index 1a9ef60..b456f0d 100644
--- a/coreml_suite/models.py
+++ b/coreml_suite/models.py
@@ -92,7 +92,10 @@ class CoreMLModelWrapper(BaseModel):
         chunked_x = chunk_batch(x, sample_shape)
         ts = list(torch.full((len(chunked_x), timestep_shape[0]), t[0]))
         chunked_context = chunk_batch(c_crossattn, context_shape)
-        chunked_control = chunk_control(control, sample_shape[0])
+
+        chunked_control = [None] * len(chunked_x)
+        if control is not None:
+            chunked_control = chunk_control(control, sample_shape[0])
 
         return chunked_x, ts, chunked_context, chunked_control