## Environment

You can install the environment using the following commands

```{bash}
conda create -n test python=3.11
pip install -r requirements.txt
```

## Training of toy problems
You can train toy problems using the following command:

```{bash}
torchrun toy_training.py --dataset vector-disjoint --dir experiments/vector_disjoint 
```
We have the other datasets in our paper which can be found in the click configuration of the `toy_training.py` file

##  MNIST

To train on MNIST simply run:

```{bash}
torchrun --nproc-per-node 1 mnist_training.py  --dir mnist_experiment
```

To sample using MNIST you can run:

```{bash}
torchrun mnist_sampling --load_checkpoint path_to_ckpt --dir mnist_samples
```