accept unets as list for decoder

2025-12-19 09:44:19 +01:00 · 2022-05-20 20:31:26 -07:00
parent f526f14d7c
commit 80497e9839
3 changed files with 5 additions and 2 deletions
--- a/train_decoder.py
+++ b/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)