(ctx, *args)
| 77 | |
| 78 | @staticmethod |
| 79 | def backward(ctx, *args): |
| 80 | if not torch.autograd._is_checkpoint_valid(): |
| 81 | raise RuntimeError("Checkpointing is not compatible with .grad(), please use .backward() if possible") |
| 82 | inputs = ctx.saved_tensors |
| 83 | # Stash the surrounding rng state, and mimic the state that was |
| 84 | # present at this time during forward. Restore the surrounding state |
| 85 | # when we're done. |
| 86 | rng_devices = [] |
| 87 | if ctx.preserve_rng_state and ctx.had_cuda_in_fwd: |
| 88 | rng_devices = ctx.fwd_gpu_devices |
| 89 | with torch.random.fork_rng(devices=rng_devices, enabled=ctx.preserve_rng_state): |
| 90 | if ctx.preserve_rng_state: |
| 91 | torch.set_rng_state(ctx.fwd_cpu_state) |
| 92 | if ctx.had_cuda_in_fwd: |
| 93 | set_device_states(ctx.fwd_gpu_devices, ctx.fwd_gpu_states) |
| 94 | detached_inputs = detach_variable(inputs) |
| 95 | with torch.enable_grad(), torch.cuda.amp.autocast(ctx.had_autocast_in_fwd): |
| 96 | outputs = ctx.run_function(*detached_inputs) |
| 97 | |
| 98 | if isinstance(outputs, torch.Tensor): |
| 99 | outputs = (outputs,) |
| 100 | torch.autograd.backward(outputs, args) |
| 101 | grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp |
| 102 | for inp in detached_inputs) |
| 103 | return (None, None) + grads |
| 104 | |
| 105 | |
| 106 | def checkpoint(function, *args, **kwargs): |
nothing calls this directly
no test coverage detected