I’m trying  transformer without teacher forcing for time series prediction and I came out with this snippet. Theoretically it works fine, but it’s tremendously slow and it keep crashing due to cuda running out of memory as soon as I use a batch size greater than 3,4 samples.
# src shape is [n_batch, src_length, d_model] --> [N, S, D]
# tgt and gt (groundtruth) shape is [n_batch, tgt_length, d_model] --> [N, T, D]

l = 0.
for src, tgt, gt in loader:
    src, tgt, gt = src.to(device), tgt.to(device), gt.to(device)
    optimizer.zero_grad()

    # first tgt element is the last src one, followed by a zero-filled tensor
    tgt = torch.cat((src[:, -1:, :], torch.zeros(tgt.shape[0], tgt.shape[1] - 1, tgt.shape[2], device=device)), dim=1)

    out = torch.zeros_like(tgt, device=device)
    memory = model.encode(src)
    for j in range(gt.shape[1]):
        out = model.decode(tgt, memory)
        l += loss(out[:, j, :], gt[:, j, :])
        tgt = tgt.detach()
        tgt = torch.cat((tgt[:, :j + 1, :], out[:, j:j+1, :], tgt[:, j + 1:gt.shape[1] - 1, :]), dim=1) 
    l.backward()