Official code for "Enabling Uncertainty Estimation in Iterative Neural Networks" (ICML 2024) - View it on GitHub