def _compile(self):
    """Compile the keras model using the correct scope."""
    # pylint: disable=protected-access
    self._init_strategy()
    assert self._strategy is not None
    with self._strategy.scope():
      if self._model_config.model == 'two_tower':
        module_model = model.TwoTowerRanker(
            self.vocab_size,
            activation=self._model_config.activation,
            matcher_name=self._model_config.matcher,
            word_emb_size=self._model_config.word_emb_size,
            hidden_size=self._model_config.hidden_size,
            dropout=self._model_config.dropout,
            use_batch_norm=self._model_config.use_batch_norm,
            contextualizer=self._model_config.contextualizer,
            context_num_layers=self._model_config.context_num_layers,
            bidirectional=self._model_config.bidirectional,
            tied_encoders=self._model_config.tied_encoders,
            embedder_name=self._model_config.embedder,
            matcher_hidden_size=self._model_config.matcher_hidden_size,
            bert_model_name=self._model_config.bert_model_name,
            bert_model_path=self._model_config.bert_model_path,
            bert_trainable=self._model_config.bert_trainable,
            bert_dropout=self._model_config.bert_dropout,
            projection_dim=self._model_config.projection_dim,
            classify_claim=self._model_config.classify_claim,
        )
        self._inner_model = module_model
        if self._encoder is None:
          raise ValueError('Encoder is not initialized.')
        # This hackery is necessary since keras doesn't handle dictionary inputs
        # well, so we have to manually specify input/output output shapes. Since
        # this is dependent on the model (e.g., bert vs other), let the encoder
        # yield this.
        inputs = self._encoder.compute_input_shapes()
        outputs = module_model(inputs)
        module_model.input_names = sorted(inputs.keys())
        module_model._feed_input_names = sorted(inputs.keys())
        module_model.output_names = sorted(
            ['evidence_matching', 'claim_classification'])
        self._model = tf.keras.Model(inputs=inputs, outputs=outputs)
        self._model.input_names = sorted(inputs.keys())
        self._model._feed_input_names = sorted(inputs.keys())
        self._model.output_names = sorted(
            ['evidence_matching', 'claim_classification'])
        self._model.summary(line_length=500)
      elif self._model_config.model == 'one_tower':
        raise NotImplementedError()
      else:
        raise ValueError('Invalid model')
      metrics = {}
      evidence_metrics = [
          tf.keras.metrics.BinaryAccuracy(name='accuracy'),
          tf.keras.metrics.Precision(name='precision'),
          tf.keras.metrics.Recall(name='recall'),
          tf.keras.metrics.AUC(name='auc'),
          tf.keras.metrics.TruePositives(name='tp'),
          tf.keras.metrics.FalsePositives(name='fp'),
          tf.keras.metrics.TrueNegatives(name='tn'),
          tf.keras.metrics.FalsePositives(name='fn'),
      ]
      metrics['evidence_matching'] = evidence_metrics

      loss = {}
      loss['evidence_matching'] = losses.WeightedBinaryCrossentropyFromProbs(
          positive_class_weight=self._model_config.positive_class_weight)

      loss_weights = {
          'evidence_matching': 1.0,
          'claim_classification': self._model_config.claim_loss_weight
      }
      if self._model_config.classify_claim:
        # TODO(perodriguez): add claim classifier metrics
        claim_metrics = [
            tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
        ]
        metrics['claim_classification'] = claim_metrics
        loss['claim_classification'] = (
            tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)
        )
      else:
        loss['claim_classification'] = losses.ZeroLoss()
        metrics['claim_classification'] = []
      self._model.compile(
          loss=loss,
          optimizer=tf.keras.optimizers.Adam(self._model_config.learning_rate),
          metrics=metrics,
          loss_weights=loss_weights,
      )
Compile the keras model using the correct scope.
language/serene/training.py