(batch_size)
| 185 | # if OOM, then halves batch_size and tries again |
| 186 | @find_executable_batch_size(starting_batch_size=self.max_batch_size) |
| 187 | def forward_batch(batch_size): |
| 188 | test_batch = torch.ones((batch_size, max_length), device=self.device).long() |
| 189 | for _ in range(5): |
| 190 | _ = F.log_softmax(self._model_call(test_batch), dim=-1).cpu() |
| 191 | return batch_size |
| 192 | |
| 193 | batch_size = forward_batch() |
| 194 | utils.clear_torch_cache() |
nothing calls this directly
no test coverage detected