(inputs)
| 384 | """See base class.""" |
| 385 | |
| 386 | def step_fn(inputs): |
| 387 | if self.config.runtime.enable_xla and (self.config.runtime.num_gpus > 0): |
| 388 | task_train_step = tf.function(self.task.train_step, jit_compile=True) |
| 389 | else: |
| 390 | task_train_step = self.task.train_step |
| 391 | logs = task_train_step( |
| 392 | inputs, |
| 393 | model=self.model, |
| 394 | optimizer=self.optimizer, |
| 395 | metrics=self.train_metrics) |
| 396 | self._train_loss.update_state(logs[self.task.loss]) |
| 397 | self.global_step.assign_add(1) |
| 398 | |
| 399 | inputs = self.next_train_inputs(iterator) |
| 400 | self.strategy.run(step_fn, args=(inputs,), options=self._runtime_options) |
nothing calls this directly
no test coverage detected