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