Function
set_modules_to_backward_prefetch
(model, num_to_backward_prefetch)
Source from the content-addressed store, hash-verified
| 24 | |
| 25 | |
| 26 | def set_modules_to_backward_prefetch(model, num_to_backward_prefetch): |
| 27 | for i, layer in enumerate(model.layers): |
| 28 | if i < num_to_backward_prefetch: |
| 29 | continue |
| 30 | layers_to_prefetch = [ |
| 31 | model.layers[i - j] for j in range(1, num_to_backward_prefetch + 1) |
| 32 | ] |
| 33 | layer.set_modules_to_backward_prefetch(layers_to_prefetch) |
| 34 | |
| 35 | |
| 36 | def main(args): |
Tested by
no test coverage detected