INNER CODE UNIT · Python

cycle

lucidrains/x-transformers · train_copy.py:19

def cycle():
    while True:
        prefix = torch.ones((BATCH_SIZE, 1)).long().to(DEVICE)
        src = torch.randint(2, NUM_TOKENS, (BATCH_SIZE, ENC_SEQ_LEN)).long().to(DEVICE)
        tgt = torch.cat((prefix, src, src), 1)
        src_mask = torch.ones(BATCH_SIZE, src.shape[1]).bool().to(DEVICE)
        yield (src, tgt, src_mask)

# instantiate model

model = XTransformer(
    dim = 128,
    tie_token_emb = True,
    return_tgt_loss = True,
    enc_num_tokens=NUM_TOKENS,
    enc_depth = 3,
    enc_heads = 8,
    enc_max_seq_len = ENC_SEQ_LEN,

View source record →