I am trying to run this program in PyTorch which is custom:

```
class NeuralNetwork(nn.Module):
def __init__(self):
super(NeuralNetwork, self).__init__()
self.flatten = nn.Flatten(start_dim=1)
self.linear_relu_stack = nn.Sequential(
nn.Linear(2142, 51),
nn.ReLU(),
nn.Linear(51, 1)
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
model = NeuralNetwork().to(device)
```

The above code is the custom NeuralNetwork. I defined the training loop below.

```
def train_loop(data, y, model, loss_fn, optimizer):
for i in range(data.shape[0]):
# Compute model prediction and loss
pred = model(data[i, :, :])
loss = loss_fn(pred, y[i, :])
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Loss: {}".format(loss.item()))
```

The following is how I would like to train it

```
final_data = torch.randn(500, 42, 51)
final_output = torch.randn(500, 1)
learning_rate = 1e-3
batch_size = 1
epochs = 5
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
epochs = 10
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train_loop(final_data, final_output, model, loss_fn, optimizer)
print('Done!!')
```

The variable `final_data`

is of shape 500x42x51. The variable `final_output`

is of shape 500x1.

I have been trying to run the above data for 10 epochs but I always end up with this error:

```
Epoch 1
-------------------------------
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
<ipython-input-17-3fb934698ecf> in <module>()
9 for t in range(epochs):
10 print(f"Epoch {t+1}\n-------------------------------")
---> 11 train_loop(final_data, final_output, model, loss_fn, optimizer)
12
13 print('Done!!')
4 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/flatten.py in forward(self, input)
38
39 def forward(self, input: Tensor) -> Tensor:
---> 40 return input.flatten(self.start_dim, self.end_dim)
41
42 def extra_repr(self) -> str:
TypeError: flatten() takes at most 1 argument (2 given)
```

The output is basically a classification between 0 or 1.

I am still a newbie in terms of PyTorch and would like some help solving this issue and understanding what’s wrong.

Thank you