def _train_squad(self, run_eagerly=False, ds_type='mirrored'):
    """Runs BERT SQuAD training. Uses mirrored strategy by default."""
    self._init_gpu_and_data_threads()
    input_meta_data = self._read_input_meta_data_from_file()
    strategy = self._get_distribution_strategy(ds_type)

    run_squad.train_squad(
        strategy=strategy,
        input_meta_data=input_meta_data,
        run_eagerly=run_eagerly,
        custom_callbacks=[self.timer_callback])
    def _train_squad(self, use_ds=True, run_eagerly=False):
        """Runs BERT SQuAD training."""
        assert tf.version.VERSION.startswith('2.')
        self._init_gpu_and_data_threads()
        input_meta_data = self._read_input_meta_data_from_file()
        strategy = self._get_distribution_strategy(use_ds)

        run_squad.train_squad(strategy=strategy,
                              input_meta_data=input_meta_data,
                              run_eagerly=run_eagerly,
                              custom_callbacks=[self.timer_callback])