From 99a0a9996d9f736236bc80dd8be316597ec27255 Mon Sep 17 00:00:00 2001 From: aszc-dev Date: Fri, 3 Nov 2023 00:31:46 +0100 Subject: [PATCH] Fix cn chunking --- coreml_suite/models.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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