add constraint plots using log_total_constraint

NSLS-II · Jan 3, 2025 · d903c68 · d903c68
1 parent 30794f1
commit d903c68
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
@@ -183,6 +183,8 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
         if obj.constraint is None:
             # test_constraint = obj.constraint_probability(model_inputs).detach().squeeze().numpy()
             test_constraint = agent.constraint(model_inputs).squeeze().numpy()
+        else:
+            test_constraint = obj.constraint_probability(model_inputs).detach().squeeze().numpy()
 
         fitness_ax = None
         fit_err_ax = None