Fix missing import

awslabs · Jun 14, 2023 · c2643a5 · c2643a5
1 parent 297ee5b
commit c2643a5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/fortuna/prob_model/posterior/name_to_posterior_state.py b/fortuna/prob_model/posterior/name_to_posterior_state.py
@@ -1,6 +1,7 @@
 import enum
 
 from fortuna.output_calib_model.state import OutputCalibState
+from fortuna.prob_model.posterior.sgmcmc.hmc.hmc_state import HMCState
 from fortuna.prob_model.posterior.laplace.laplace_state import LaplaceState
 from fortuna.prob_model.posterior.map.map_state import MAPState
 from fortuna.prob_model.posterior.normalizing_flow.advi.advi_state import ADVIState
@@ -21,3 +22,4 @@ class NameToPosteriorState(enum.Enum):
     vars()[SWAGState.__name__] = SWAGState
     vars()[SGHMCState.__name__] = SGHMCState
     vars()[CyclicalSGLDState.__name__] = CyclicalSGLDState
+    vars()[HMCState.__name__] = HMCState