How can I get the batch shape in a Model class?

This question is more like basic python Class question I think.(If this question is not belong here, plz let me know)
I am trying to build a simple conditional GAN model.
I created my own custom Dataset class and collate_fn which is fed to DataLoader class.
I customized those since input data has varying shape.
Thankfully, I followed this solution using rnn.pad_sequence function in order to padding within each batches.(thx to @jdhao and @pinocchio)

n_epochs = 100

for epoch in range(n_epochs):

    for batch_idx, (data_input, labels) in enumerate(data_loader):


Above code gives me the output looks like the following one.

torch.Size([10, 86, 408])
torch.Size([10, 105, 408])
torch.Size([10, 110, 408])
torch.Size([10, 87, 408])
torch.Size([10, 76, 408])
torch.Size([10, 44, 408])
torch.Size([10, 54, 408])
torch.Size([10, 52, 408])
torch.Size([10, 89, 408])
torch.Size([10, 59, 408])
torch.Size([10, 60, 408])
torch.Size([10, 43, 408])

So you can see that first one is the batch length, but the second varying number represents zero-padded input.

The problem is that I need to get these second number inside Generator class.

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

        self.label_emb = nn.Embedding(n_class, n_class)

        self.model = nn.Sequential(
            nn.Linear(latent_dim + n_class, 256),
            nn.LeakyReLU(0.2, inplace= True),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2, inplace= True),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2, inplace= True),
            nn.Linear(1024, int(dt_h * dt_w)),

    def forward(self, z, labels):

        gen_input =, z), -1)

        tech = self.model(gen_input)

        tech = tech.view(batch_size, dt_h, dt_w)
        return tech

The dt_h in the code is the one that should be replaced by the varying number.

Any advice will be appriciated. Thanks! :fist_right: :fist_left:

if batch_size and dt_w are of constant size, you can just replaced dt_h with -1 and it will automatically select the right size.

1 Like

First, Thanks to the advice. I will give it a try.
Just one thing in mind keep bothering me.
Shouldn’t I write the code something like self.data_input and get the number with self.data_input.size()[1] inside the Generator class? :thinking:

I get the idea of your advice since .view() is act like reshape in numpy.
However What about the dt_h in the last layer before the sigmoid activation?

You can indeed do something like self.data_input.shape[1] to get the same on-the-fly if you’d want to. The -1 (as you said) is just to exploit reshaping and get pytorch to automatically select how many rows I can get if I specify the other dimensions.

With regards to dt_h in the last nn.Linear module, I’m not too sure how to resolve that. I know there’s a nn.LazyLinear which takes arbitrary inputs but there’s no equivalent for outputs (documentation here). You could do something along the lines of defining the final nn.Linear as 1024 inputs and the output as 110 * 408 (as 110 is the maximum output I see from your example output above), then you can just pad all outputs with zeros so they have the same shape.

Perhaps, on the forum somewhere someone has asked about linear layers with variable outputs?

Thanks for all the answers! you are the best! :hugs: