This dataset is a HF mirror of the official challenge training data for this challenge:
https://www.codabench.org/competitions/8934/
To ease the split along the nuisance parameter axis, the dataset challenge has been reordered as
ncosmo, np, ... -> np, ncosmo, ...
dset = datasets.load_dataset("cosmostat/neurips-wl-challenge")
dset = dset.with_format('torch')
