parameters: to_dot: Allow exporting single-node trees

1 files changed, 5 insertions, 1 deletions

diff --git a/lib/parameters.py b/lib/parameters.py
index cb4b76f..01b3e66 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -704,7 +704,11 @@ class ModelAttribute:
         return ret
 
     def to_dot(self):
-        if type(self.model_function) == df.SplitFunction:
+        if type(self.model_function) in (
+            df.SplitFunction,
+            df.StaticFunction,
+            df.AnalyticFunction,
+        ):
             import pydot
 
             graph = pydot.Dot("Regression Model Tree", graph_type="graph")