Searched refs:_maybe_compute_loss (Results 1 – 1 of 1) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/distributed/pipelining/ |
| H A D | schedules.py | 369 def _maybe_compute_loss(self, stage, output, target_mbs, mb_index): member in _PipelineSchedule
714 self._maybe_compute_loss(self._stage, output, target_mbs, i)
815 self._maybe_compute_loss(self._stage, output, target_mbs, fwd_mb_index)
852 self._maybe_compute_loss(self._stage, output, target_mbs, fwd_mb_index)
1261 self._maybe_compute_loss(stage, output, target_mbs, mb_index)
1562 self._maybe_compute_loss(stage, output, target_mbs, mb_index)
|