Similar to torch.cuda.amp.autocast, but compatible with torch 1.5.1
(enabled: bool)
| 399 | |
| 400 | @contextmanager |
| 401 | def autocast(enabled: bool) -> Generator: |
| 402 | """Similar to torch.cuda.amp.autocast, but compatible with torch 1.5.1""" |
| 403 | if enabled: |
| 404 | with torch.cuda.amp.autocast(enabled): |
| 405 | yield |
| 406 | else: |
| 407 | yield |
| 408 | |
| 409 | |
| 410 | class CheckpointFunction(torch.autograd.Function): |