Extend part of the model

Hello,

I am wondering is it possible to extend part of the model, if I only want to get the batch gradient of the last several layers?

I think `model = extend(model)` will waste memory if only the batch gradient of the last several layers is needed.

For example, if I only want to extend the last two layers (let's say the last two layers are fc1 and fc2) of a large model, can I do something like this:

```
model.fc1 = extend(model.fc1)
model.fc2 = extend(model.fc2)
```