Defined two nn.ModuleList() in similar way but only one works fine

ShaRinGan: <p>I’m building a multi-task-learning network (Segmentation and Depth) and for that I choose U-Net Architecture. I used common encoder , separate bottle-neck and decoder. <strong>The issue I found after debugging is; I have defined two nn.ModuleList() , one for decoder_seg and another for decoder_depth; in both of them I have added layers ; but while looping through layers in forward() method, it’s only looping through decoder_seg but not through decoder_depth (it’s length is showing 0), although both are build in similar way</strong></p>
<p>Here’s the code:</p>
<pre><code class="lang-auto">
import torch
import torch.nn as nn
import torchvision.transforms.functional as F
import torch.optim


class IntermediateBlocks(nn.Module):
    def __init__(self, block_in_channels, block_out_channels):
        super(IntermediateBlocks, self).__init__()
        self.block = nn.Sequential(

            nn.Conv2d(block_in_channels, block_out_channels,
                      kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(block_out_channels),
            nn.ReLU(inplace=True),

            nn.Conv2d(block_out_channels, block_out_channels,
                      kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(block_out_channels),
            nn.ReLU(inplace=True),

        )

    def forward(self, x):
        return self.block(x)


class DepthSegmentation(nn.Module):
    def __init__(self, in_channels=3, out_channels=3, intermediate_channels=None):
        super(DepthSegmentation, self).__init__()
        self.out_channels = out_channels

        if intermediate_channels is None:
            intermediate_channels = [64, 128, 256, 512]

        """ ---------------- Down-Sampling Layers --------------- """
        self.encoder = nn.ModuleList()
        for num_channels in intermediate_channels:
            self.encoder.append(IntermediateBlocks(block_in_channels=in_channels, block_out_channels=num_channels))
            in_channels = num_channels
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        """ ---------------- Bottle Neck Layers ------------------ """
        # One for Segmentation
        self.bottleneck_seg = IntermediateBlocks(intermediate_channels[-1], intermediate_channels[-1] * 2)
        # one for Depth Estimation
        self.bottleneck_depth = IntermediateBlocks(intermediate_channels[-1], intermediate_channels[-1] * 2)

        """ ----------------- Up-Sampling Layers ---------------- """
        self.decoder_intermediate_channels = reversed(intermediate_channels)  # [512, 256, 128, 64]

        # for segmentation
        self.decoder_seg = nn.ModuleList()
        for num_channels in self.decoder_intermediate_channels:
            self.decoder_seg.append(
                nn.ConvTranspose2d(
                    num_channels * 2, num_channels, kernel_size=2, stride=2
                )
            )
            self.decoder_seg.append(IntermediateBlocks(num_channels * 2, num_channels))

        # for depth estimation
        self.decoder_depth = nn.ModuleList()
        for num_channels in self.decoder_intermediate_channels:
            self.decoder_depth.append(
                nn.ConvTranspose2d(
                    num_channels * 2, num_channels, kernel_size=2, stride=2
                )
            )
            self.decoder_depth.append(IntermediateBlocks(num_channels * 2, num_channels))

        """ -------------------- Final Layers ---------------------"""
        self.final_seg = nn.Conv2d(
            in_channels=intermediate_channels[0], out_channels=self.out_channels,
            kernel_size=1, stride=1, padding=0)
        self.final_depth = nn.Conv2d(
            in_channels=intermediate_channels[0], out_channels=self.out_channels,
            kernel_size=1, stride=1, padding=0)


    def forward(self, x):
        skip_connections_layers = []

        for layers in self.encoder:
            # First, processing it through Intermediate Block consisting of few conv layers
            x = layers(x)

            """ Since, the encoder was made from few IntermediateBlocks,--- 
            so for getting the skip_connections_layers (for concatenation in Decoder part),
            --- we are going to append the last layer of each IntermediateBlocks """
            skip_connections_layers.append(x)  # here the x supplied is from last layer of each Intermediate Block

            # MaxPooling is applied after every Intermediate Blocks (for Down-Sampling)
            x = self.pool(x)

        common_encoder_seg_output = x
        common_encoder_depth_output = x

        seg_out = self.bottleneck_seg(common_encoder_seg_output)
        depth_out = self.bottleneck_depth(common_encoder_depth_output)

        # As, every up-sampled layer need to be concatenated with last element(layer) present in
        # skip_connections_layers list so, it's better to reverse the list
        skip_connections_layers = skip_connections_layers[::-1]

        """ Since, self.decoder_seg is like ==&gt;
        # [convTranspose2D, IntermediateBlocks, convTranspose2D, IntermediateBlocks, ..... ]
        # Here, the concatenation will be happening with only convTranspose2D layers, therefore
        # while looping, we have to use step = 2 """

        print(f'Length of decoder_depth: {len(self.decoder_depth)}')

        print(f'Length of decoder_seg: {len(self.decoder_seg)}',"\n")

        for i_seg in range(0, len(self.decoder_seg), 2):
            print(f"seg_out - Before {i_seg // 2} convTranspose2D: {seg_out.shape}")
            # First processing with convTranspose2D
            seg_out = self.decoder_seg[i_seg](seg_out)
            print(f"seg_out - After {i_seg // 2} convTranspose2D: {seg_out.shape}")
            required_skip_layer = skip_connections_layers[i_seg // 2]

            # While concatenation, we set dim = 1, because we want to do it along depth(channels)
            # Since a batch consist of  ==&gt; (batch_size, channels_dim, height, width)
            # Also, we need to make sure the shape matches
            if seg_out.shape != required_skip_layer.shape:
                # [2:]  ==&gt;  height, width
                seg_out = F.resize(seg_out, size=required_skip_layer.shape[2:])
            concatenated_layer_seg = torch.cat((required_skip_layer, seg_out), dim=1)

            print(f"seg_out - Before {i_seg // 2} IntermediateBlocks: {seg_out.shape}")
            # After that, processing with IntermediateBlocks
            seg_out = self.decoder_seg[i_seg + 1](concatenated_layer_seg)
            print(f"seg_out - After {i_seg // 2} IntermediateBlocks: {seg_out.shape}", '\n')


        print(f'seg_out - Before passing to final layer: {seg_out.shape}')

        print("\n",f'Length of decoder_depth: {len(self.decoder_depth)}')


        # Similarly for Depth Estimation head
        for i_depth in range(0, len(self.decoder_depth), 2):
            print(f"Before depth_out: {depth_out.shape}")
            # First processing with convTranspose2D
            depth_out = self.decoder_depth[i_depth](depth_out)
            print(f"After depth_out: {depth_out.shape}")
            required_skip_layer = skip_connections_layers[i_depth // 2]

            if depth_out.shape != required_skip_layer.shape:
                # [2:]  ==&gt; height, width
                depth_out = F.resize(depth_out, size=required_skip_layer.shape[2:])
            concatenated_layer_depth = torch.cat((required_skip_layer, depth_out), dim=1)

            # After that, processing with IntermediateBlocks
            depth_out = self.decoder_depth[i_depth + 1](concatenated_layer_depth)
            print(f"After concatenated_layer_depth: {depth_out.shape}")


        return self.final_seg(seg_out), self.final_depth(depth_out)



# Dummy Input
input_batch = torch.randn((16, 3, 160, 160))
model = DepthSegmentation(in_channels=3, out_channels=3, intermediate_channels=[64,128,256,512])
seg_output, depth_output = model(input_batch)

print(seg_output.shape)
print(depth_output.shape)
</code></pre>
<h2><a name="this-is-the-error-i-got-1" class="anchor" href="#this-is-the-error-i-got-1"></a>This is the error I got :</h2>
<p>RuntimeError: Given groups=1, weight of size [3, 64, 1, 1], expected input[16, 1024, 10, 10] to have 64 channels, but got 1024 channels instead.</p>
<p><strong>This is full error with debugging output:</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/a/3a0c5f4adf6beb87643eace459a78f33d73e45a7.png" data-download-href="https://discuss.pytorch.org/uploads/default/3a0c5f4adf6beb87643eace459a78f33d73e45a7" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/a/3a0c5f4adf6beb87643eace459a78f33d73e45a7_2_690x474.png" alt="image" data-base62-sha1="8hwcrpDtUxTyuctHt6Tw7vyjagT" width="690" height="474" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/a/3a0c5f4adf6beb87643eace459a78f33d73e45a7_2_690x474.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/a/3a0c5f4adf6beb87643eace459a78f33d73e45a7_2_1035x711.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/a/3a0c5f4adf6beb87643eace459a78f33d73e45a7.png 2x" data-dominant-color="242529"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1270×873 130 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ShaRinGan: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Please have a look.</p>

ptrblck: <p><code>self.final_depth(depth_out)</code> fails since <code>depth_out</code> has a shape of <code>[16, 1024, 10, 10]</code> while <code>self.final_depth</code> expects an activation with 64 channels.</p>