reset_classifier: fix it to actually work, expose toggle param for gpu #7
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently the reset_classifier function does not work, as there is no self.out_dim in the model.
I've updated reset_classifier to work by using the current head in_features to determine the in channels for the new head.
I also removed global_pool param, which was dead weight and was not being used.
I added a force option in case you are not changing num_classes but still want to reset (defaults to off).
Finally, added a to_gpu flag b/c in building a model that is on cpu atm, you don't want to force only the head to gpu if the entire model will be moved later, which was how the initial version was setup.