The size of tensor a (64) must match the size of tensor b (7) at non-singleton dimension 1

lambchop97: <p>I am using an LSTM model to train a model on the Toronto emotional speech set (TESS)<br>
dataset <a href="https://www.kaggle.com/datasets/ejlok1/toronto-emotional-speech-set-tess" rel="noopener nofollow ugc">here</a>. My code for the model and the training loop is as follow:</p>
<pre><code class="lang-auto">class TorchLSTMNet(torch.nn.Module):
    def __init__(self, inputs=1, outputs=7, bidirectional=False):
        super().__init__()

        self.lstm = torch.nn.LSTM(inputs, 128, batch_first=True, bidirectional=bidirectional)
        self.linear = torch.nn.Sequential(
            torch.nn.Linear(128 * 2 if bidirectional else 128, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, outputs)
        )
        
    def forward(self, x):
        # x shape: (samples, steps, inputs)
        out, (h, c) = self.lstm(x)
        
        # use final output
        if h.shape[0] == 2:
            return self.linear(torch.cat([h[0], h[1]], axis=1))
        else:
            return self.linear(h[0])

net = TorchLSTMNet()
print(net)
</code></pre>
<pre><code class="lang-auto">def train(net, train_iter, test_iter, num_epochs, lr, device = d2l.try_gpu()):
    """Train a model with a GPU (defined in Chapter 6)."""
    print('training on', device)
    net.to(device)
    print(net)
    optimizer = torch.optim.SGD(net.parameters(), lr=lr)
    loss = torch.nn.CrossEntropyLoss()
    animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
                            legend=['train loss', 'train acc', 'test acc'])
    timer, num_batches = d2l.Timer(), len(train_iter)
    for epoch in range(num_epochs):
        # Sum of training loss, sum of training accuracy, no. of examples
        metric = d2l.Accumulator(3)
        net.train()
        for i, (X, y) in enumerate(train_iter):
            timer.start()
            optimizer.zero_grad()
            X, y = X.to(device), y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            l.backward()
            optimizer.step()
            with torch.no_grad():
                metric.add(l * X.shape[0], d2l.accuracy(y_hat, y), X.shape[0])
            timer.stop()
            train_l = metric[0] / metric[2]
            train_acc = metric[1] / metric[2]
            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
                animator.add(epoch + (i + 1) / num_batches,
                             (train_l, train_acc, None))
        test_acc = d2l.evaluate_accuracy_gpu(net, test_iter)
        animator.add(epoch + 1, (None, None, test_acc))
    print(f'loss {train_l:.3f}, train acc {train_acc:.3f}, '
          f'test acc {test_acc:.3f}')
    print(f'{metric[2] * num_epochs / timer.sum():.1f} examples/sec '
          f'on {str(device)}')
</code></pre>
<p>The data is of shape <code>torch.Size([64, 40, 1])</code> and <code> torch.Size([64, 7])</code> for each batch of <code>x</code> and <code>y</code>, respectively. <code>X</code> is a tensor of size <code>(5600, 40, 1)</code> which contains the 40 mfcc features, and <code>Y</code> is a tensor of size <code>(5600, 7)</code> which contains the associated emotions.</p>
<pre><code class="lang-auto">data_loaders = {
    'train': torch.utils.data.DataLoader(train, shuffle=True, batch_size=64),
    'test':   torch.utils.data.DataLoader(test, batch_size=64),
}

# print the x and y shapes for one minibatch
for (x, y) in data_loaders['train']:
    print(x.shape, y.shape)
    break
</code></pre>
<p>The code fails at the training loop, at line <code>metric.add(l * X.shape[0], d2l.accuracy(y_hat, y), X.shape[0])</code>, with the error in the title.</p>
<p>I would really appreciate any help.</p>