Pytorch: Inference on Specific Layers of a Model

Question

I’m building a flexible module that implements Monte Carlo Dropout for Bayesian Inference on any kind of model. (Implementing MC dropout by adding dropout to the fc layers d