Hi,
This most likely happens because the first part of your model doesn’t get gradient because of some quirks of how checkpointing works.
Can you try making data['data'] require gradients before giving it to the checkpoint? (you can ignore the computed gradient, just add a data['data'].requires_grad_() if it is a Tensor.