add stage_task in the diagram

AllenNeuralDynamics · Dec 21, 2023 · 82a8338 · 82a8338
1 parent 65c2b14
commit 82a8338
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/code/aind_auto_train/plot/curriculum.py b/code/aind_auto_train/plot/curriculum.py
@@ -43,7 +43,7 @@ def draw_diagram_rules(curriculum):
     # Add nodes (stages)
     for stage in stages:
         dot.node(name=stage.name,
-                 label=stage.name,
+                 label=f'{stage.name}\n{curriculum.parameters[stage].task}',
                  shape='ellipse',
                  style='filled',
                  fillcolor=stage_color_mapper[stage.name],