From 6447afa4c2013949f806a243b5758c2d7c4fbbd7 Mon Sep 17 00:00:00 2001 From: Hu Ye Date: Fri, 24 Nov 2023 20:19:04 +0800 Subject: [PATCH] fix a bug the original code don't register the downsamper module into the parent Module. --- library/original_unet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/original_unet.py b/library/original_unet.py index 240b8595..23e60786 100644 --- a/library/original_unet.py +++ b/library/original_unet.py @@ -485,7 +485,7 @@ class DownBlock2D(nn.Module): self.resnets = nn.ModuleList(resnets) if add_downsample: - self.downsamplers = [Downsample2D(out_channels, out_channels=out_channels)] + self.downsamplers = nn.ModuleList([Downsample2D(out_channels, out_channels=out_channels)]) else: self.downsamplers = None