Checkpoint wrapper doesn't support dict return type

## 🐛 Bug

I have been trying to wrap some modules which return dict output type with checkpoint_wrapper. It looks like it only returns keys back.

## To Reproduce

```
from fairscale.nn import checkpoint_wrapper
from torch import nn

class A(nn.Module):
def forward(self, x):
return {"x": 1, "y": 2}

a = checkpoint_wrapper(A())
print(a(None))
## Prints out ("x", "y")
```

## Expected behavior

I should get dict back as it is.

## Environment

fairscale version: 0.3.7