INNER CODE UNIT · Python

DEC_SEQ_LEN

lucidrains/x-transformers · train_copy.py:14

DEC_SEQ_LEN = 64 + 1
DEVICE = default_device()

# helpers

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,

View source record →