Thanks @ptrblck !
I believe that in the example that you created it doesn’t throw any errors because you are not checkpointing the first layer, as you first run self.features on the input, and then only checkpoint an identity mapping lambda,
which kinda defeats the purpose.
Perhaps your intention was to create a dummy layer before the feature extraction ? This is what I’m trying to do right now.
This is my attempt of a workaround, now in standalone code thanks to your example : )
If you have an idea for making a dummy layer that is more “light weight” and/or some code base solution (that makes checkpoint function ignore the fact that the input has no gradients) I’ll appreciate it a lot 
import torch
from torch import nn
from torchvision import models
from torch.utils.checkpoint import checkpoint

class DummyLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.dummy = nn.Parameter(torch.ones(1, dtype=torch.float32))
    def forward(self,x):
        return x + self.dummy - self.dummy #(also tried x+self.dummy)

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)

        self.dummy_layer = DummyLayer()

    def forward(self, x):
        x = self.dummy_layer(x)
        x = checkpoint(self.features, x)
        print(x.shape)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

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

While it works in this synthetic example (and using these “small” images), in my actual scenario, I’m on the very edge of gpu memory even for a single slice (and I need all of it for proper feature extraction),
and even this dummy layer makes me get OOM (even if I don’t use checkpointing)