Fix cn chunking

2023-11-01 01:08:14 +01:00
parent d63df5b62f
commit dfdc1bf520
1 changed files with 1 additions and 1 deletions
@@ -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