Support new dataset field in loader, instead of old train field

[stefan-apollo]

Previously mod add training data info in config.yaml was stored in train, now it is stored in dataset. The code is not updated yet, but the existing config files are also not updated. Thus the issue only occurs in newly trained models.

Todo: Implement a fix that works for new and old files. Not hard but also just not urgent.

Code in branch fix/load_new_config_keys

[danbraunai-apollo]

loader.py in main now looks like this:

        modulus = modulus or cfg["dataset"]["modulus"]
        fn_name = fn_name or cfg["dataset"]["fn_name"]
        if frac_train is not None:
            assert frac_train == cfg["dataset"]["frac_train"], (
                f"frac_train provided ({frac_train}) does not match frac_train in tlens model "
                f"config ({cfg['dataset']['frac_train']})"
            )
        else:
            frac_train = cfg["dataset"]["frac_train"]
        seed = seed if seed is not None else cfg["dataset"]["seed"]

Are you satisfied with this @stefan-apollo? It won't work for old models, like your description asked for, but it does approximately match what your branch did.

Fwiw, in this case, I'd prefer not to add extra code inside the config to handle older models, but rather have custom code in scripts that use older models which converts things to make it readable with the latest codebase. But lmk if you think this might be problematic or non optimal.