Fix appending a new block to input_blocks

This commit is contained in:
Subuday 2024-02-15 08:55:52 +00:00
parent fd6c0afbbf
commit 8676ab30d9
1 changed files with 3 additions and 1 deletions

View File

@ -76,7 +76,7 @@ class UNet(nn.Module):
)
self.input_blocks = nn.ModuleList([])
block_in_channels = in_channels
block_in_channels = in_channels * 2
for _ in range(num_blocks):
block = nn.ModuleList([])
@ -88,6 +88,8 @@ class UNet(nn.Module):
)
)
self.input_blocks.append(block)
self.middle_blocks = nn.ModuleList([])
self.output_blocks = nn.ModuleList([])