Passing arguments to forward method

Hi All,
I am instantiating an architecture like this:

class Discriminator_classwise(nn.Module):
    def __init__(self, inc=4096, num_class=126):
        super(Discriminator, self).__init__()
        self.classfier_list = []
        for _ in range(num_class):
            self.classfier_list.append(nn.Sequential(nn.Linear(inc,512),nn.Linear(512,512),nn.Linear(512,2)))

    def forward(self, x, reverse=False, eta=1.0, choose_class = 0):
        if reverse:
            x = grad_reverse(x, eta)
        which_classifier = self.classfier_list[choose_class]
        x_out = which_classifier(x)
        return x_out

Essentially, what I am doing is I want to choose a classifier from the list of classifier according to the class that is predicted for my input.

Could someone help me in writing a forward function for this, which takes an argument which is a class index and returns the forward passed tensor through the desired classifier?

Thanks in advance,
Megh

Hi,

You simply need to make your list a ModuleList so that it is tracked properly:

self.classfier_list = nn.ModuleList()

And then the code you shared will work just fine :slight_smile:

Thanks a lot @albanD. But I am wondering, why will a normal list not work? Anyways, I will just access the appropriate network from the list that is it, right?

Does this make sense? Please let me know if something is unclear.

A regular list doesn’t work because the nn construct are not opening plain list for speed reasons.
But the ModuleList will behave just like a regular python list :slight_smile: