GPU Memory usage

Hi, thanks for your great work! I wonder if the GPU memory usage will grow linearly with the parameter "num_ens" in the code? If I increase "num_ens", there will be an OOM error definitely, right?  Maybe it's better to implement the MC sampling process in the "forward" function of the "BBBConv" and "BBBLinear" to avoid this issue.