From dfdc1bf5201edde37a05a98bec6e3f40761eca1f Mon Sep 17 00:00:00 2001 From: aszc-dev Date: Wed, 1 Nov 2023 01:08:14 +0100 Subject: [PATCH] Fix cn chunking --- coreml_suite/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coreml_suite/models.py b/coreml_suite/models.py index b047056..64eba09 100644 --- a/coreml_suite/models.py +++ b/coreml_suite/models.py @@ -97,5 +97,5 @@ 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 = chunk_control(control, x.shape[0]) return chunked_x, ts, chunked_context, chunked_control