Decoder only transformer model

1

I am trying to run an ordinary differential equation within decoder only transformer model.
My ultimate aim is to plot loss and training curves of the model upon reversing tokenization.
However, I came across following error during training my model.
Error: IndexError : index out of range
I could not have figured out what went wrong although tried to modify the code many times.
Much appreciated any suggestion or help. Thanks.
Below is the code:

import torch
import torch.optim as optim
from x_transformers import AutoregressiveWrapper, TransformerWrapper, Decoder

n_bins = 20
n_tokens = n_bins * n_bins

def convert_to_tokens(Y, n_bins):
digit_x = np.digitize(np.log(x), np.linspace(-3, 3, n_bins))
digit_y = np.digitize(np.log(y), np.linspace(-3, 3, n_bins))
return digit_y*n_bins + digit_y

model = TransformerWrapper(
num_tokens = n_tokens,
max_seq_len = 1024,
attn_layers = Decoder(
dim = 512,
depth = 6,
heads = 8
)
)

optimizer = optim.Adam(model.parameters(), lr=1e-3)

from scipy import integrate
for epoch in range(20):

optimizer.zero_grad
for i in range(100):
    a, b, c, d = 0.4, 0.002, 0.001, 0.7
    def f(Y, t):
        x, y = Y
        return [a * x - b * x * y,
        c * x * y - d * y]

    xy0 = [5, 4]
    t = np.linspace(0, 10, 250)
    Y = integrate.odeint(f, xy0, t)
    inputs = convert_to_tokens(Y, n_bins)
    inputs = torch.tensor(inputs).unsqueeze(0)

    logits = model(inputs)
    loss = torch.nn.functional.cross_entropy(logits, inputs)
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item()
    if i % 5 == 0:print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 20)) 
    running_loss = 0.0