Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-09-22 06:44:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow alternative functions to load extra data.

Browse Source 
PiperOrigin-RevId: 434495304
This commit is contained in:
Sven Gowal
2022-05-26 17:43:03 +01:00
committed by alimuldal
parent 3cb13ea0b9
commit b2fa23c838
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
adversarial_robustness/jax/experiment.py
+3 -2
View File
@@ -480,8 +480,9 @@ class Experiment(experiment.AbstractExperiment):
def _extra_train_dataset(self) -> tfds.typing.Tree[np.ndarray]:
"""Creates the training dataset."""
load_fn = (datasets.load_dummy_data if self.config.dry_run else
datasets.load_extra)
load_extra_fn = (self.config.training.get('load_extra_fn', None) or
datasets.load_extra)
load_fn = datasets.load_dummy_data if self.config.dry_run else load_extra_fn
load_fn = functools.partial(
load_fn, path_npz=self.config.training.extra_data_path)
ds = _dataset(