INNER CODE UNIT · Python
length
lucidrains/x-transformers · train_discover_reversal.py:121
length = torch.randint(1, len_max + 1, (1,)).item()
src = torch.randint(0, num_fillers, (batch_size, length), device = device)
tgt = src.flip(-1)
logits = model(src, tgt)
loss = torch.nn.functional.cross_entropy(logits.reshape(-1, num_tokens), tgt.reshape(-1))
loss.backward()
optimizer.step()
optimizer.zero_grad()
if (i + 1) % 1000 == 0:
accs = evaluate()
if sum(accs.values()) > best_acc:
best_acc = sum(accs.values())
best_state = {name: param.detach().clone() for name, param in model.named_parameters()}