Hi,
I am considering the use of gradient checkpointing to lessen the VRAM load. From what I understand there were some issues with stochastic nodes (e.g. Dropout) at some point in time to apply gradient checkpointing.
However I have a kind of Bayesian Neural Network which needs quite a bit of memory, hence I am interested in gradient checkpointing. Since all the weights are bayesian, there is stochasticity everywhere which can’t be removed.
Can I safely apply gradient checkpointing to this Bayesian Neural Network ?