Hey!
I’m building a flexible module that implement Monte Carlo Dropout for Bayesian Inference on any kind of model.
(Implementing MC dropout by adding dropout to fc layers during test time and inference for multiple times to get distribution of outputs)
To save time, I want the to divide the model into two parts:

No dropout implemented so run only once
Layers which are after the first dropout (first fc layer) therefore will run multiple times.

I can’t use nn.Sequential() because many of the models I work with have functionality in forward (like x.permute()).
Do you have any Idea how to inference on only specific layers without changing the model itself or the forward() code?