fix bug caused by taking length of itertools.pairwise() (#288)

pinellolab · Jun 6, 2024 · 74dc742 · 74dc742
1 parent 6530de4
commit 74dc742
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/dnadiffusion/models/unet.py b/src/dnadiffusion/models/unet.py
@@ -33,7 +33,7 @@ def __init__(
         self.init_conv = nn.Conv2d(input_channels, init_dim, (7, 7), padding=3)
         dims = [init_dim, *(dim * m for m in dim_mults)]
 
-        in_out = itertools.pairwise(dims)
+        in_out = list(zip(dims[:-1], dims[1:]))
         block_klass = partial(ResnetBlock, groups=resnet_block_groups)
 
         # time embeddings