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:
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: