@ptrblck thanks for your response!
Yes, this model is just part of a larger network, i.e., top_vec which is the output of this model is being used by another model. I see top_vec as a vector that has the encoded version of vector x (i…e, src) by the BERT. I a sense, the weights associated with this class should be updated (i.e., learned) during training. Upon my investigations, I noticed that this part of the model consumes much of a memory, so that I thought it’d be better to checkpoint it. Isn’t that true argument, wherever we have learning (i.e., the update of model parameters), we can use gradient checkpointing??

I also noticed that there’s a recently implemented option in Huggingface’s BERT which allows us to apply gradient checkpointing easily. That’s an argument that is specified in BertConfig and then the object is passed to   BertModel.from_pretrained. I also tried that, but have the same above issues that I mentioned: 1) the performance does not yield to that of setting without gradient-checkpointing. 2) the training is much faster. While I expect it to be slower, in expense of having memory usage reduced, as Gradient Checkpointing makes the model recompute activations wherever necessary instead of retrieving them, which totally results in more computations, and more training time.