Searched defs:_prepare_forward_infra (Results 1 – 2 of 2) sorted by relevance
| /external/pytorch/test/distributed/pipelining/ | ||
| D | test_schedule.py | 53 def _prepare_forward_infra(self, n_microbatches): member in MockPipelineStage |
| /external/pytorch/torch/distributed/pipelining/ | ||
| D | stage.py | 254 def _prepare_forward_infra(self, num_microbatches: int): member in _PipelineStageBase |