flaxmodels
flaxmodels copied to clipboard
Remove head option in Resnet
It doesn’t look like there’s a way to remove the head in the resnet models. There is a way to dump the activations. It would be nice if I could specify up till which resnet block I’d like outputs for to avoid wasted computation of the stages that I don’t need (if I just want to use it as a low level feature extractor)