Accidentally removed OneHot and ReduceArgmax from subterm

unda-ml · Apr 4, 2024 · d77efc0 · d77efc0
1 parent a3e2c7c
commit d77efc0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/core/graph/subterm.rs b/src/core/graph/subterm.rs
@@ -65,6 +65,7 @@ impl Context {
                         }
                     Operation::ReduceMax { node, dim: _ }
                     | Operation::ReduceMean { node, dim: _ }
+                    | Operation::ReduceArgmax { node, dim: _ }
                     | Operation::ReduceSum { node, dim: _ } => {
                         to_visit.push(node);
                     }
@@ -73,6 +74,7 @@ impl Context {
                         to_visit.push(on_true);
                         to_visit.push(on_false);
                     }
+                    Operation::OneHot(node) => to_visit.push(node),
                     Operation::Constant(_) | Operation::Parameter(_) => {}
                 }
                 node_map.insert(self.nodes[node_id].clone(), node_id);