albanD:

If you put the whole network into it, you might need to set requires_grad=True on the input though. Does it help?


Thanks. Yes, it does help. Although I’m not totally happy with the result either (setting requires_grad=true to a tensor that doesn’t need it):
if y.requires_grad is False:
   y.requires_grad_(True)  # require gradient for y even when learning net, otherwise checkpoint doesn't work
   result = torch.utils.checkpoint.checkpoint(self.net, x, y)