class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.LSTM(
            input_size=6,
            hidden_size=6,
            num_layers=2,
            batch_first=True,
        )

    def forward(self, x):
        out, (h_n, h_c) = self.rnn(x, None)
        return out[:, -1, :]    # Return output at last time-step

X = torch.FloatTensor(X)
y = torch.LongTensor(Y)

rnn = RNN()
optimizer = torch.optim.Adam(rnn.parameters(), lr=0.001)
loss_func = nn.CrossEntropyLoss()     


for j in range(500):
    for i, item in enumerate(X):
        item = item.unsqueeze(0)
        print(item)
        print(item.shape)
        output = rnn(item)
        target = y[i]
        print(target)
        target = target.squeeze_()
        print(output,target)
        loss = loss_func(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

ValueError: Expected input batch_size (1) to match target batch_size (2).
output after forward = tensor([[0.1039, 0.0309, 0.1265, 0.1670, 0.0287, 0.1056]], grad_fn=)
target = tensor([0, 0])