From b2fa23c8381ec3deb01d06430920a3ad9aec4471 Mon Sep 17 00:00:00 2001 From: Sven Gowal Date: Mon, 14 Mar 2022 17:06:39 +0000 Subject: [PATCH] Allow alternative functions to load extra data. PiperOrigin-RevId: 434495304 --- adversarial_robustness/jax/experiment.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/adversarial_robustness/jax/experiment.py b/adversarial_robustness/jax/experiment.py index 20c037f..3b87faf 100644 --- a/adversarial_robustness/jax/experiment.py +++ b/adversarial_robustness/jax/experiment.py @@ -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(