Fairscale checkpoint-wrapper deconstructs NamedTuple outputs

It seems `checkpoint_wrapper` deconstructs NamedTuple into just a Tuple. This makes it difficult to do type checks for interfaces between models.

Here is a sample test to reproduce:

```
from typing import NamedTuple
from fairscale.nn import checkpoint_wrapper
from torch import Tensor
from torch.nn import (
GELU,
Conv2d,
Flatten,
Linear,
MaxPool2d,
Module,
Sequential,
)

class SimpleModelOutput(NamedTuple):
intermidiate: Tensor
final: Tensor

class SimpleModel(Module):
def __init__(
self,
output_dim: int,
):
super().__init__()

self.output_dim = output_dim
self.model_1 = Sequential(
Conv2d(3, 20, kernel_size=5),
GELU(),
MaxPool2d(2, stride=2),
Conv2d(20, 50, kernel_size=5),
GELU(),
)

self.model_2 = Sequential(
MaxPool2d(2, stride=2),
Flatten(),
Linear(50 * 4 * 4, output_dim),
GELU(),
)

def forward(self, x: Tensor) -> SimpleModelOutput:
intermidiate: Tensor = self.model_1(x)
final: Tensor = self.model_2(intermidiate)

return SimpleModelOutput(intermidiate=intermidiate, final=final)

def test_return_type_for_checkpoint_wrapper():
model_1 = SimpleModel(500)

input_data = torch.rand(((1, 3, 28, 28)))
out1 = model_1(input_data)

model_2 = checkpoint_wrapper(model_1)
out2 = model_2(input_data)

# checkpoint_wrapper seems to deconstruct the output type of this model into a tuple
assert type(out1) == type(out2), f'Non-matching types: {type(out1)=} != {type(out2)=}'
```

Here is the output:


Start with the checkpoint_wrapper entry point and run the reproduction test from the issue using the SimpleModelOutput NamedTuple. Trace how the wrapped model returns its output; done means the wrapped and unwrapped models return the same NamedTuple type and the assertion passes.

Written by the indexing model from the issue text.