Here is what i do:
class Resnet50Extractor(nn.Module):
    def __init__(self, submodule, extracted_layer):
        super(Resnet50Extractor, self).__init__()
        self.submodule = submodule
        self.extracted_layer = extracted_layer

    def forward(self, x):
        if self.extracted_layer == 'maxpool':                       
            modules = list(self.submodule.children())[:4]
        elif self.extracted_layer == 'inner-layer-3':                     
            modules = list(self.submodule.children())[:6]
            third_module = list(self.submodule.children())[6]
            third_module_modules = list(third_module.children())[:3]    # take the first three inner modules
            third_module = nn.Sequential(*third_module_modules)
            modules.append(third_module)
        elif self.extracted_layer == 'layer-3':                     
            modules = list(self.submodule.children())[:7]
        else:                                               # after avg-pool
            modules = list(self.submodule.children())[:9]

        self.submodule = nn.Sequential(*modules)
        x = self.submodule(x)
        return x

then call like below:
model_ft = models.resnet50(pretrained=True)
extractor = Resnet50Extractor(model_ft, extracted_layer)
features = extractor(input_tensor)