I am trying to implement a Dropout layer using pytorch as follows:

```
class DropoutLayer(nn.Module):
def __init__(self, p):
super().__init__()
self.p = p
def forward(self, input):
if self.training:
u1 = (np.random.rand(*input.shape)<self.p) / self.p
u1 *= u1
return u1
else:
input *= self.p
```

And then calling a simple NN.sequential:

```
model = nn.Sequential(nn.Linear(input_size,num_classes), DropoutLayer(.7), nn.Flatten())
opt = torch.optim.Adam(model.parameters(), lr=0.005)
train(model, opt, 5) #train(model, optimizer, epochs #)
```

But I’m getting the following error:

```
TypeError: flatten() takes at most 1 argument (2 given)
```

Not sure what I’m doing wrong. Still new to pytorch. Thanks.