Thanks for the explanation. I misunderstood the issue.
I played around a bit and also noticed, that the problem occurs when you checkpoint the first operation.
However, could you explain your use case a bit more?
You have a model which calculates features for several images.
I assume each image has an own “computation path”. Is this right?
Each of this submodules are checkpointed to save some memory and somehow you get the gradient error from your first post.
Did you freeze the model and just use it for the feature extraction or do you train it?
Could you post a small dummy model showing your use case?