diff --git a/official/recommendation/ncf_keras_benchmark.py b/official/recommendation/ncf_keras_benchmark.py index 010c4a196c9..ac1870c4215 100644 --- a/official/recommendation/ncf_keras_benchmark.py +++ b/official/recommendation/ncf_keras_benchmark.py @@ -113,6 +113,11 @@ def benchmark_1_gpu(self): self._setup() self._run_and_report_benchmark() + def benchmark_2_gpus(self): + self._setup() + FLAGS.num_gpus = 2 + self._run_and_report_benchmark() + class KerasNCFSyntheticData(KerasNCFBenchmarkBase): """Benchmark NCF model using synthetic data.""" @@ -149,3 +154,8 @@ def _extract_benchmark_report_extras(self, stats): def benchmark_1_gpu(self): self._setup() self._run_and_report_benchmark() + + def benchmark_2_gpus(self): + self._setup() + FLAGS.num_gpus = 2 + self._run_and_report_benchmark()