Hi @Rishabh_Dahale,
I don’t think it is possible.
But if your problem is that the computation graph is too big to fit in your GPU memory, gradient checkpointing might be able to help you here.
Please take  a look at the official doc 11 and at this tutorial 10. Even though it is using the Variable API, the checkpointing is still relevant.