We need to not train on the whole 2M sample (too long to experiment). Note that these samples are unshuffled and we need to shuffle them as there is a big block of signal and a big block of bkg.
We should convert all instances of U[-1, 1] noise to N(0, 1) noise
We should make some plotting code to evaluate after training. Maybe start with trying to make some plots for the average jets that look like the ones in the paper. There is some code for plotting average / individual jets here.
In the arXiv paper, we say that we tried different things along the way (FCN, CNN, 2-stream, etc.) and that although we don't report the results for all, the architectures are available on github. We should add them to the models/ folder.