def benchmark_dqn_8_gpu_xla(self):
utils.set_session_config(enable_xla=True)
strategy = distribution_strategy_utils.get_distribution_strategy(
distribution_strategy='default', num_gpus=8)
self._run(strategy, batch_size=64 * 8)
def benchmark_dqn_no_dist_strat_1_gpu_xla(self):
utils.set_session_config(enable_xla=True)
strategy = distribution_strategy_utils.get_distribution_strategy(
distribution_strategy='off', num_gpus=1)
self._run(strategy)
