Thanks for the clarification and the nicely simplified code!
Maybe we could add a dummy layer between the self.feature_extractor and the append.
I created a small example with a resnet18, which at least doesn’t throw any errors.
Could you try that?
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.features = nn.Sequential(*list(models.resnet18(pretrained=False).children())[:5])
        self.fc1 = nn.Linear(200704, 2)
        
    def forward(self, x):
        x = self.features(x)
        x = checkpoint(lambda x: x, x)
        print(x.shape)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

model = MyModel()
x = torch.randn(1, 3, 224, 224)
output = model(x)
output.mean().backward()
print(model.features[0].weight.grad)