(function, *args, **kwargs)
| 46 | |
| 47 | |
| 48 | def checkpoint(function, *args, **kwargs): |
| 49 | if flags.get_checkpointing(): |
| 50 | kwargs.setdefault("use_reentrant", True) |
| 51 | return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) |
| 52 | else: |
| 53 | return function(*args, **kwargs) |
| 54 | |
| 55 | |
| 56 | def downscale_pos(pos): |