Merge pull request #5 from abarankab/dev

nn.Identity() not having parentheses bug fix
abarankab · Aug 19, 2021 · 579e3b2 · 579e3b2
2 parents 9da9534 + 423d72d
commit 579e3b2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ddpm/unet.py b/ddpm/unet.py
@@ -297,7 +297,7 @@ def __init__(
                               time_emb_dim=time_emb_dim, num_classes=num_classes, activation=activation, norm=norm, num_groups=num_groups),
                 ResidualBlock(out_channels, out_channels, dropout,
                               time_emb_dim=time_emb_dim, num_classes=num_classes, activation=activation, norm=norm, num_groups=num_groups),
-                AttentionBlock(out_channels) if use_attn else nn.Identity,
+                AttentionBlock(out_channels) if use_attn else nn.Identity(),
                 Downsample(out_channels, use_reflection_pad=use_reflection_pad) if not is_last else nn.Identity(),
             ]))
 
@@ -316,7 +316,7 @@ def __init__(
                               time_emb_dim=time_emb_dim, num_classes=num_classes, activation=activation, norm=norm, num_groups=num_groups),
                 ResidualBlock(in_channels, in_channels, dropout,
                               time_emb_dim=time_emb_dim, num_classes=num_classes, activation=activation, norm=norm, num_groups=num_groups),
-                AttentionBlock(in_channels) if use_attn else nn.Identity,
+                AttentionBlock(in_channels) if use_attn else nn.Identity(),
                 Upsample(in_channels, align_corners=align_corners, use_reflection_pad=use_reflection_pad),
             ]))