--gradient_accumulation_steps 1 --n_gpus 2 --patch_size 16 --dim 128 --out_dim 768 --channels 128 64 32 16 --depth 12 --num_heads 16 --mlp_dim 1024 --dropout 0.2 --patch_dropout 0.2 --latent_channels 4 --m2c_modality both --m2c_vision_dropout 0.1 --random_air_layers --random_roll --m2c_trainable_pe --m2c_d_ff 1024 --m2c_latent_dim 256 --m2c_out_dim 768 --m2c_num_vision_heads 8 --m2c_groups 16 --m2c_num_heads 16 --m2c_num_layers 16 --m2c_dropout 0.2 --batch_size 128 --data_path ../datasets/simple_structures --max_seq_len  8 --use_tensorboard --save_path ./ICLR_FINAL2_checkpoints/A_small --VAE_reconstruction_loss BCE --lr 1e-4 --epochs 500 --warmup_epochs 50 --beta_clip 1.0 --beta_kl 1e-3 --queue_size 131072 --weight_decay 2e-4
--gradient_accumulation_steps 1 --n_gpus 2 --patch_size 16 --dim 128 --out_dim 768 --channels 128 64 32 16 --depth 12 --num_heads 16 --mlp_dim 1024 --dropout 0.2 --patch_dropout 0.2 --latent_channels 4 --m2c_modality both --m2c_vision_dropout 0.1 --random_air_layers --random_roll --m2c_trainable_pe --m2c_d_ff 1024 --m2c_latent_dim 256 --m2c_out_dim 768 --m2c_num_vision_heads 8 --m2c_groups 16 --m2c_num_heads 16 --m2c_num_layers 16 --m2c_dropout 0.2 --batch_size 128 --data_path ../datasets/normal_structures --max_seq_len 12 --use_tensorboard --save_path ./ICLR_FINAL2_checkpoints/A_large --VAE_reconstruction_loss BCE --lr 1e-4 --epochs 500 --warmup_epochs 50 --beta_clip 1.0 --beta_kl 1e-3 --queue_size 131072 --weight_decay 2e-4 --use_amp