Add square-symmetry augmentations and sampling fixes#8

JiamingPan merged 2 commits into

mainnkern/cosmo_diffusion:mainfrom

codex/augmentationsnkern/cosmo_diffusion:codex/augmentationsCopy head branch name to clipboard