python main.py --dataset mnist --epochs 100
python main.py --dataset fashionmnist
python main.py --dataset svhn --n_channels 3 --image_size 32 --embed_dim 128
python main.py --dataset cifar10 --n_channels 3 --image_size 32 --embed_dim 128

