sweep_config = {
    'pretrained_path': "path/to/phase/one",
    # 'pretrained_path': None,
    'retake_training': False,
    'dim_feedforward': 1024,
    'nheads': 16,
    'masking_p': 0.15,
    'nlayers': 12,
    'dropout': 0.0,
    'dim_model': 512,
    'batch_first': True,
    'n_tokens': 20340, 
    'batch_size': 64, # 64 for phase one, 6 for phase two 
    'context_length': 1500,
    'lr': 1e-4,
    'decay':1e-1,
    'warmup': 5000, # 5000 for phase one, 0 for phase two
    'max_epochs': 1226446, 
    'autoregressive': False,
    'pool': None,
    'supervised_task': False,
    'learnable_pe': True,
    'spatial_aware': False, # False for phase one, True for phase two
    'data_path': "path/to/output"
    }
