reformer-crime and punishment-original config

config = {
    "attention_head_size": 64,
    "attn_layers": ["local", "lsh", "local", "lsh", "local", "lsh"],
    "axial_pos_embds": True,
    "sinusoidal_pos_embds": False,
    "axial_pos_embds_dim": [64, 192],
    "axial_pos_shape": [512, 1024],
    "lsh_attn_chunk_length": 64,
    "local_attn_chunk_length": 64,
    "feed_forward_size": 512,
    "hidden_act": "relu",
    "hidden_size": 256,
    "is_decoder": True,
    "max_position_embeddings": 524288,
    "num_attention_heads": 2,
    "num_buckets": [64, 128],
    "num_hashes": 1,
    "vocab_size": 320,
    "lsh_attention_probs_dropout_prob": 0.0,
    "lsh_num_chunks_before": 1,
    "lsh_num_chunks_after": 0,
    "local_num_chunks_before": 1,
    "local_num_chunks_after": 0,
    "local_attention_probs_dropout_prob": 0.025,
    "hidden_dropout_prob": 0.025,
}