INNER CODE UNIT · Python
best_acc
lucidrains/x-transformers · train_discover_reversal.py:115
best_acc = -1.
best_state = None
for i in range(train_steps):
model.train()
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: