How to exclude some operations in checkpoint wrapper?

I would like to checkpoint my module that takes the result of a checkpointed module (`cond` in the example below) as input.

```python
class Test(nn.Module):
def __init__(self):
super(Test, self).__init__()

def forward(x, cond=None):
if cond is not None:
# do something

return result
```

The above module works fine when checkpointed and `cond` is None. However, when `cond` is not None, I am getting the following error.

```
RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed).
Saved intermediate values of the graph are freed when you call .backward() or autograd.grad().
Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.
```

I have tried the following context managers but none of them works. This is a large module and I prefer to wrap it while ignoring the statements that cannot be checkpointed. Please kindly advise any workarounds. Thank you.

```python
from fairscale.nn.checkpoint.checkpoint_activations import disable_checkpointing

with disable_checkpointing():
if cond is not None:
# do something
```

```python
from fairscale.nn.checkpoint.checkpoint_activations import enable_recomputing

with enable_recomputing():
if cond is not None:
# do something
```

Reproduce the failure with a non-None cond in the checkpointed Test module, then read fairscale.nn.checkpoint.checkpoint_activations around disable_checkpointing and enable_recomputing. Determine whether either context manager is intended to exclude operations from recomputation; done means a validated workaround or a clearly documented limitation.

Written by the indexing model from the issue text.