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
```