Model.forward doesn't work while features and labels have the same size for the first dimension of shape

Mona_Jalal: <p>I am not sure what accidental change in my code caused this. Could you please guide me how to fix this?</p>
<p>The code is:</p>
<pre><code class="lang-auto">    if train:
        print('training...')
        torch.autograd.set_detect_anomaly(True)
        for i_batch, sample_batched in enumerate(dataloader_train):
            feats = torch.stack(sample_batched['image'])
            labels = torch.as_tensor(sample_batched['label']).cuda()
            print('feats shape: ', feats.shape)
            print('labels shape: ', labels.shape)
            pred, labels, loss = model.forward(feats, labels)
            output = model(feats)
            loss = criterion(output, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            acc = (output.argmax(dim=1) == labels).float().mean()
            train_preds = output.argmax(dim=1)
</code></pre>
<pre><code class="lang-auto">feats shape:  torch.Size([64, 419, 512])
labels shape:  torch.Size([64])
Traceback (most recent call last):
  File "main_classifier.py", line 289, in &lt;module&gt;
    pred, labels, loss = model.forward(feats, labels)
  File "/home/jalal/research/venv/dpcc/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 166, in forward
    return self.module(*inputs[0], **kwargs[0])
  File "/home/jalal/research/venv/dpcc/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() takes 2 positional arguments but 3 were given
</code></pre>

Mona_Jalal: <aside class="quote no-group" data-username="Mona_Jalal" data-post="1" data-topic="145141">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mona_jalal/48/25715_2.png" class="avatar"> Mona_Jalal:</div>
<blockquote>
<pre><code class="lang-auto">if train:
        print('training...')
        torch.autograd.set_detect_anomaly(True)
        for i_batch, sample_batched in enumerate(dataloader_train):
            feats = torch.stack(sample_batched['image'])
            labels = torch.as_tensor(sample_batched['label']).cuda()
            print('feats shape: ', feats.shape)
            print('labels shape: ', labels.shape)
            pred, labels, loss = model.forward(feats, labels)
            output = model(feats)
            loss = criterion(output, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            acc = (output.argmax(dim=1) == labels).float().mean()
            train_preds = output.argmax(dim=1)
</code></pre>
</blockquote>
</aside>
<p>Problem was uncommenting a line from the past. I removed “model.forward” line. Could you please confirm if the code below is correct? Thanks a lot.</p>
<pre data-code-wrap="if"><code class="lang-if">        print('training...')
        torch.autograd.set_detect_anomaly(True)
        for i_batch, sample_batched in enumerate(dataloader_train):  
              
           
            feats = torch.stack(sample_batched['image']) 
            labels = torch.as_tensor(sample_batched['label']).cuda() 
            print('feats shape: ', feats.shape)
            print('labels shape: ', labels.shape)
            ##pred, labels, loss = model.forward(feats, labels)
            output = model(feats)
            loss = criterion(output, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            acc = (output.argmax(dim=1) == labels).float().mean()
            train_preds = output.argmax(dim=1)
</code></pre>