accept unets as list for decoder

dalle2_pytorch/dalle2_pytorch.py

@@ -59,6 +59,9 @@ def default(val, d):
     return d() if isfunction(d) else d
 
 def cast_tuple(val, length = 1):
+    if isinstance(val, list):
+        val = tuple(val)
+
     return val if isinstance(val, tuple) else ((val,) * length)
 
 def module_device(module):

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.3.6',
+  version = '0.3.7',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',

train_decoder.py

@@ -86,7 +86,7 @@ def create_decoder(device, decoder_config, unets_config):
         ))
     
     decoder = Decoder(
-        unet=tuple(unets),  # Must be tuple because of cast_tuple
+        unet=unets,
         **decoder_config
     )
     decoder.to(device=device)