So, as a simple example, I would say this should work, but it does not output the same as the target after a lot of epochs. With ADam it almost doesn’t move, with SGD or Adam with high LR,  it converges to a wrong solution where each row is filled with entries 50-65. CrossEntropyLoss seems to be for multiclass classification, so thats why I tried MSELoss.
Any ideas how to make output converge to the target?
And how to sample multiple batches from the series?
import torch
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

data = torch.arange(0,1000000,1).float()

feature_size = 16
sequence_size = 100



src = torch.as_strided(data,(sequence_size,feature_size),(1,1)).unsqueeze(1)
target = torch.as_strided(data[1:],(sequence_size,feature_size),(1,1)).unsqueeze(1)
print(target.size())
# size = (Sequence, Batch, Features)
tf_model= nn.Transformer(feature_size,8,2,2,2,0.2)

src_mask = tf_model.generate_square_subsequent_mask(sequence_size)
optimizer = torch.optim.SGD(tf_model.parameters(),lr=0.1)
loss_fn = torch.nn.MSELoss()

for epoch in range(1000):
    out = tf_model(src,target,src_mask)
    optimizer.zero_grad()
    loss = loss_fn(out,target)
    loss.backward()
    optimizer.step()

out = tf_model(src,target)
print(target)
print(out)
#Should be the same after some training?