diffcharge update

DAI-Lab · Dec 29, 2024 · 8e5e6cc · 8e5e6cc
1 parent b0d4f17
commit 8e5e6cc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/generator/diffcharge/diffusion.py b/generator/diffcharge/diffusion.py
@@ -94,6 +94,7 @@ def __init__(self, cfg: DictConfig):
         self.beta_start = cfg.model.beta_start
         self.beta_end = cfg.model.beta_end
         self.current_epoch = 0
+        self.conditioning_var_n_categories = cfg.dataset.conditioning_vars
 
         self.conditioning_module = ConditioningModule(
             categorical_dims=cfg.dataset.conditioning_vars,