Internal change

PiperOrigin-RevId: 411096556

Internal change
PiperOrigin-RevId: 411096556
eaa4003c · Abdullah Rashwan · A. Unique TensorFlower · 19e3d234 · eaa4003c
Commit eaa4003c authored Nov 19, 2021 by Abdullah Rashwan Committed by A. Unique TensorFlower Nov 19, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

official/core/base_trainer.py official/core/base_trainer.py +4 -1

No files found.
--- a/official/core/base_trainer.py
+++ b/official/core/base_trainer.py
@@ -403,7 +403,10 @@ class Trainer(_AsyncTrainer):
    """See base class."""

    def step_fn(inputs):
-      task_train_step = self.task.train_step
+      if self.config.runtime.enable_xla and (self.config.runtime.num_gpus > 0):
+        task_train_step = tf.function(self.task.train_step, jit_compile=True)
+      else:
+        task_train_step = self.task.train_step
      logs = task_train_step(
          inputs,
          model=self.model,