(ctx, run_function, preserve_rng_state, *args)
| 56 | |
| 57 | @staticmethod |
| 58 | def forward(ctx, run_function, preserve_rng_state, *args): |
| 59 | check_backward_validity(args) |
| 60 | ctx.run_function = run_function |
| 61 | ctx.preserve_rng_state = preserve_rng_state |
| 62 | ctx.had_autocast_in_fwd = torch.is_autocast_enabled() |
| 63 | if preserve_rng_state: |
| 64 | ctx.fwd_cpu_state = torch.get_rng_state() |
| 65 | # Don't eagerly initialize the cuda context by accident. |
| 66 | # (If the user intends that the context is initialized later, within their |
| 67 | # run_function, we SHOULD actually stash the cuda state here. Unfortunately, |
| 68 | # we have no way to anticipate this will happen before we run the function.) |
| 69 | ctx.had_cuda_in_fwd = False |
| 70 | if torch.cuda._initialized: |
| 71 | ctx.had_cuda_in_fwd = True |
| 72 | ctx.fwd_gpu_devices, ctx.fwd_gpu_states = get_device_states(*args) |
| 73 | ctx.save_for_backward(*args) |
| 74 | with torch.no_grad(): |
| 75 | outputs = run_function(*args) |
| 76 | return outputs |
| 77 | |
| 78 | @staticmethod |
| 79 | def backward(ctx, *args): |
nothing calls this directly
no test coverage detected