Pytorch Weight Initialization problem for DCGAN

I am trying to create a generator for DCGAN and initialize custom weights. In the Pytorch tutorial, the code is given as below:

# Generator Code

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)

However, I am trying to create blocks for sequential operations and and my generator looks as follows:

class ConvTr(nn.Module):
    def __init__(self, input_channels, output_channels, k_size, stride, pad, b = False):
        super().__init__()
        
        self.conv = nn.ConvTranspose2d(in_channels=input_channels,
                                       out_channels=output_channels,
                                       kernel_size=k_size,
                                       stride=stride,
                                       padding=pad,
                                       bias=b)
        
        self.batch_norm = nn.BatchNorm2d(input_channels)
        self.activation = nn.ReLU()
        
        
    def forward(self, x):
        x = self.conv(x)
        x = self.activation(x)
        x = self.batch_norm(x)
        return x



class Generator(nn.Module):
    
    def __init__(self, inputNoise=100):
        super().__init__()

        self.conv1 = ConvTr(input_channels = inputNoise, output_channels = 128, k_size = 4, stride = 2, pad = 0)
        self.conv2 = ConvTr(input_channels = 128, output_channels = 64, k_size = 3, stride = 2, pad = 1)
        self.conv3 = ConvTr(input_channels = 64, output_channels = 32, k_size = 4, stride = 2, pad = 1)
        self.conv4 = ConvTranspose2d(in_channels=32, out_channels=1, kernel_size=4, stride=2, padding=1)
        self.tanh = Tanh()
        

    def forward(self, x):
        
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        output = self.tanh(x)
        return output

Here is the weight initialization function:

def weights_init(model):
    # get the class name
    classname = model.__class__.__name__
    # check if the classname contains the word "conv"
    if classname.find("Conv") != -1:
        # intialize the weights from normal distribution
        nn.init.normal_(model.weight.data, 0.0, 0.02)
    # otherwise, check if the name contains the word "BatcnNorm"
    elif classname.find("BatchNorm") != -1:
        # intialize the weights from normal distribution and set the
        # bias to 0
        nn.init.normal_(model.weight.data, 1.0, 0.02)
        nn.init.constant_(model.bias.data, 0)

Now when I am trying to initialize the weights, it is giving me the following error.

raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'ConvTr' object has no attribute 'weight'

three problems:

1. use model.apply to do module level operations (like init weight)
2. use isinstance to find out what layer it is
3. do not use .data, it has been deprecated for a long time and should always be avoided whenever possible

to initialize the weight, do the following

def init_weight(model):
    def _init_weight(module):
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
            nn.init.normal_(module.weight, 0., 0.02)
        elif isinstance(module, nn.BatchNorm2d):
            nn.init.normal_(module.weight, 1.0, 0.02)
            nn.init.constant_(module.bias, 0)

    model.apply(_init_weight)

# m is the model you want to initialize
init_weight(m)

edit: added ConvTranspose in condition