fix for auto splitting of args and kwargs into imagen, for #301

lucidrains · Jan 13, 2023 · bd303fb · bd303fb
1 parent 8e0e97c
commit bd303fb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/imagen_pytorch/trainer.py b/imagen_pytorch/trainer.py
@@ -178,7 +178,7 @@ def split_args_and_kwargs(*args, split_size = None, **kwargs):
     split_kwargs_index = len_all_args - dict_len
 
     split_all_args = [split(arg, split_size = split_size) if exists(arg) and isinstance(arg, (torch.Tensor, Iterable)) else ((arg,) * num_chunks) for arg in all_args]
-    chunk_sizes = tuple(map(len, split_all_args[0]))
+    chunk_sizes = tuple(map(len, first_tensor.split(split_size, dim = 0)))
 
     for (chunk_size, *chunked_all_args) in tuple(zip(chunk_sizes, *split_all_args)):
         chunked_args, chunked_kwargs_values = chunked_all_args[:split_kwargs_index], chunked_all_args[split_kwargs_index:]

diff --git a/imagen_pytorch/version.py b/imagen_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.18.7'
+__version__ = '1.18.9'