PyTorch TypeError: flatten() takes at most 1 argument (2 given)
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(&quot;Loss: {}&quot;.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&quot;Epoch {t+1}\n-------------------------------&quot;)
    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)
&lt;ipython-input-17-3fb934698ecf&gt; in &lt;module&gt;()
      9 for t in range(epochs):
     10     print(f&quot;Epoch {t+1}\n-------------------------------&quot;)
---&gt; 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) -&gt; Tensor:
---&gt; 40         return input.flatten(self.start_dim, self.end_dim)
     41 
     42     def extra_repr(self) -&gt; 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

The error is a bit misleading, if you try running the code from a fresh kernel, the issues are elsewhere...
There are multiple issues with your code:

You are not using the correct shape for the target tensor y, here it should have a single dimension since the output tensor is two-dimensional.

The target tensor should be of dtype Long

When iterating over you data and selecting input (and target) with data[i, :, :] (and y[i, :]), you are essentially removing the batch axis. However all builtin nn.Module work with a batch axis. You can do a slice to avoid that side effect: with data[i:i+1] and y[i:i+1] respectively. Also do note that x[j, :, :] is identical to x[j].


That being said, the usage of the cross-entropy loss is not justified. You are outputting a single logit, so it doesn't make sense to use a cross-entropy loss.

You can either output two logits on the last layer of your model,

or switch to another loss function such as a binary cross-entropy loss (either using nn.BCELoss or a nn.BCEWithLogitsLoss which includes a sigmoid activation).
In this case the target vector should be of dtype float, and its shape should equal that of pred.



def train_loop(data, y, model, loss_fn, optimizer):
    for i in range(data.shape[0]):
        pred = model(data[i:i+1])
        loss = loss_fn(pred, y[i:i+1].float())
        # [...]

final_data = torch.randn(500, 42, 51)
final_output = torch.randint(0, 2, (500,1))
loss_fn = nn.BCEWithLogitsLoss()

Then the following will work:
&gt;&gt;&gt; train_loop(final_data, final_output, model, loss_fn, optimizer)

