I tried to refactor a linear layer that takes bn and activation functions as optional arguments to be used in a NN, the code for it is
act_fn = nn.ReLU(inplace=True)

def linear_layer(n_in, n_out, drop=False, dropout_rate = 0.5, bn = False, zero_bn = False, act=True):
    lin = nn.Linear(n_in, n_out)
    dropout = nn.Dropout(p=dropout_rate)
    bn = nn.BatchNorm1d(n_out)
    nn.init.constant_(bn.weight, 0. if zero_bn else 1.)

    layers = [lin(n_in, n_out)]
    if bn:
        layers.append(bn)
    if drop:
        layers.append(dropout)
    if act: 
        layers.append(act_fn)

    return nn.Sequential(*layers)

But when I tried to test it
# test linear layer

tmp = torch.randn(64*7*7)
a = linear_layer(64*7*7, 1024)
print(a(tmp).shape)

I get
---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

<ipython-input-61-1c1d696ee893> in <module>()
      1 
      2 tmp = torch.randn(64*7*7)
----> 3 a = linear_layer(64*7*7, 1024)
      4 print(a(tmp).shape)

1 frames

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    539             result = self._slow_forward(*input, **kwargs)
    540         else:
--> 541             result = self.forward(*input, **kwargs)
    542         for hook in self._forward_hooks.values():
    543             hook_result = hook(self, input, result)

TypeError: forward() takes 2 positional arguments but 3 were given

However, it works fine if I define each layer separately and call each one in the forward() of a nn module
I even tried to reduce it down to
def linear_layer(n_in, n_out):
    lin = nn.Linear(n_in, n_out)

    layers = [lin(n_in, n_out)]


    return nn.Sequential(*layers)


But it still gives me the same error