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