In the following code, I've encountered an issue where the loss during training does not go down, please identify the error and fix it.

# -*- coding: utf-8 -*-
"""NLP_Assignment_2.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1VWepVR3oD2WVzqIRJz9wDZvQNyujeOCW
"""

import numpy as np
import pandas as pd
import torch
import torch.nn
import torch.nn.utils.rnn
import torch.utils.data
import matplotlib.pyplot as plt
import seaborn as sns
# import opencc
import os
from sklearn.model_selection import train_test_split

df_train = pd.read_csv('arithmetic_train.csv')
df_eval = pd.read_csv('arithmetic_eval.csv')
df_train = df_train.iloc[:,1:]
df_eval = df_eval.iloc[:,1:]
df_train.head()

id_to_char = {0:'0', 1:'1', 2:'2', 3:'3', 4:'4', 5:'5', 6:'6', 7:'7', 8:'8', 9:'9', 10:'+', 11:'-', 12:'*', 13:'(', 14:')', 15:'=', 16:'<pad>', 17:'<eos>'}
char_to_id = {v: k for k, v in id_to_char.items()}

vocab_size = len(char_to_id)
print('Vocab size{}'.format(vocab_size))

# transform the input data to string
df_train['tgt'] = df_train['tgt'].apply(lambda x: str(x))
df_train['src'] = df_train['src'].add(df_train['tgt'])
df_train['len'] = df_train['src'].apply(lambda x: len(x))

df_eval['tgt'] = df_eval['tgt'].apply(lambda x: str(x))
df_eval['src'] = df_eval['src'].add(df_eval['tgt'])
df_eval['len'] = df_eval['src'].apply(lambda x: len(x))

# Transform src and tgt columns to sequences of IDs
eos = [char_to_id['<eos>']]
pad = [char_to_id['<pad>']]

df_train['char_id_list'] = df_train['src'].apply(lambda x: [char_to_id[c] for c in x] + eos)
df_eval['char_id_list'] = df_eval['src'].apply(lambda x: [char_to_id[c] for c in x] + eos)

labels = []
for i in range(len(df_train['tgt'])):
    exp = df_train['tgt'][i]
    padding = pad * int(df_train['len'][i]-len(exp))
    labels.append(padding + [char_to_id[c] for c in exp] + eos)
df_train.insert(4, "label_id_list", labels)

labels = []
for i in range(len(df_eval['tgt'])):
    exp = df_eval['tgt'][i]
    padding = pad*int(df_eval['len'][i]-len(exp))
    labels.append(padding + [char_to_id[c] for c in exp] + eos)
df_eval.insert(4, "label_id_list", labels)

df_train.head()

class Dataset(torch.utils.data.Dataset):
    def __init__(self, sequences):
        self.sequences = sequences

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, index):
        x = self.sequences['char_id_list'][index][:-1]
        y = self.sequences['label_id_list'][index][1:]
        return x, y

# collate function, used to build dataloader
def collate_fn(batch):
    batch_x = [torch.tensor(data[0]) for data in batch]
    batch_y = [torch.tensor(data[1]) for data in batch]
    batch_x_lens = torch.LongTensor([len(x) for x in batch_x])
    batch_y_lens = torch.LongTensor([len(y) for y in batch_y])
    # max_len_x = batch_x_lens.max().item()

    # Pad the input sequence
    pad_batch_x = torch.nn.utils.rnn.pad_sequence(batch_x,
                                                  batch_first=True,
                                                  padding_value=char_to_id['<pad>'])

    pad_batch_y = torch.nn.utils.rnn.pad_sequence(batch_y,
                                                  batch_first=True,
                                                  padding_value=char_to_id['<pad>'])

    # pad_batch_y = torch.cat([pad_batch_y, torch.full((pad_batch_y.size(0), max_len_x - pad_batch_y.size(1)),
    #                             char_to_id['<pad>'], dtype=torch.long)], dim=1)

    return pad_batch_x, pad_batch_y, batch_x_lens, batch_y_lens

batch_size = 256
epochs = 2
embed_dim = 256
hidden_dim = 256
lr = 0.001
grad_clip = 1

ds_train = Dataset(df_train[['char_id_list', 'label_id_list']])
ds_eval = Dataset(df_eval[['char_id_list', 'label_id_list']])

dl_train = torch.utils.data.DataLoader(ds_train,
                                       batch_size=batch_size,
                                       shuffle=True,
                                       collate_fn=collate_fn)
dl_eval = torch.utils.data.DataLoader(ds_eval,
                                      batch_size=batch_size,
                                      shuffle=False,
                                      collate_fn=collate_fn)

class CharRNN(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super(CharRNN, self).__init__()

        self.embedding = torch.nn.Embedding(num_embeddings=vocab_size,
                                            embedding_dim=embed_dim,
                                            padding_idx=char_to_id['<pad>'])

        self.rnn_layer1 = torch.nn.LSTM(input_size=embed_dim,
                                        hidden_size=hidden_dim,
                                        batch_first=True)

        self.rnn_layer2 = torch.nn.LSTM(input_size=hidden_dim,
                                        hidden_size=hidden_dim,
                                        batch_first=True)

        self.linear = torch.nn.Sequential(torch.nn.Linear(in_features=hidden_dim,
                                                          out_features=hidden_dim),
                                          torch.nn.ReLU(),
                                          torch.nn.Linear(in_features=hidden_dim,
                                                          out_features=vocab_size))

    def forward(self, batch_x, batch_x_lens):
        return self.encoder(batch_x, batch_x_lens)

    # The forward pass of the model
    def encoder(self, batch_x, batch_x_lens):
        batch_x = self.embedding(batch_x)

        batch_x = torch.nn.utils.rnn.pack_padded_sequence(batch_x,
                                                          batch_x_lens,
                                                          batch_first=True,
                                                          enforce_sorted=False)

        batch_x, _ = self.rnn_layer1(batch_x)
        batch_x, _ = self.rnn_layer2(batch_x)

        batch_x, _ = torch.nn.utils.rnn.pad_packed_sequence(batch_x,
                                                            batch_first=True)

        batch_x = self.linear(batch_x)

        return batch_x

    def generator(self, start_char, max_len=200):

        char_list = [char_to_id[c] for c in start_char]

        next_char = None

        while len(char_list) < max_len:
            x = torch.tensor([char_list]).to(device)

            # Inference through the network
            x = self.embedding(x)
            x, _ = self.rnn_layer1(x)
            x, _ = self.rnn_layer2(x)
            y = self.linear(x)

            next_char = y.argmax(dim=1).item()

            if next_char == char_to_id['<eos>']:
                break

            char_list.append(next_char)

        return [id_to_char[ch_id] for ch_id in char_list]

torch.manual_seed(2)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CharRNN(vocab_size,
                embed_dim,
                hidden_dim)

criterion = torch.nn.CrossEntropyLoss(ignore_index=char_to_id['<pad>'])
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

from tqdm import tqdm
from copy import deepcopy
model = model.to(device)
model.train()
i = 0
for epoch in range(1, epochs+1):
    model.train()
    bar = tqdm(dl_train, desc=f"Train epoch {epoch}")
    for batch_x, batch_y, batch_x_lens, batch_y_lens in bar:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        optimizer.zero_grad()

        batch_pred_y = model(batch_x, batch_x_lens)

        print(batch_pred_y)
        print('='*10)
        print(batch_y)
        # Reshape batch_pred_y and batch_y to fit CrossEntropyLoss requirements
        batch_pred_y = batch_pred_y.view(-1, vocab_size)  # Flatten to (batch_size * seq_len, vocab_size)
        batch_y = batch_y.view(-1)  # Flatten to (batch_size * seq_len)

        # Calculate the loss
        loss = criterion(batch_pred_y, batch_y)
        loss.backward()

        torch.nn.utils.clip_grad_value_(model.parameters(), grad_clip) # Gradient clipping

        optimizer.step()

        i += 1
        if i % 50 == 0:
            bar.set_postfix(loss=loss.item())

    # Evaluation
    bar = tqdm(dl_eval, desc=f"Validation epoch {epoch}")
    matched, total = 0, 0
    model.eval()
    with torch.no_grad():
        for batch_x, batch_y, batch_x_lens, batch_y_lens in bar:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)

            # Get predictions
            predictions = model(batch_x, batch_x_lens)
            # predictions = predictions.argmax(dim=1)
            # print(predictions)
            # Calculate exact match (EM)
            matched += (predictions.argmax(dim=2) == batch_y).all(dim=1).sum().item()
            total += len(batch_y)

    print("EM = ", matched/total)