add battery transition models to representers

ahalev · Apr 5, 2023 · f2e113f · f2e113f
1 parent 91dc011
commit f2e113f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/pymgrid/utils/serialize.py b/src/pymgrid/utils/serialize.py
@@ -20,6 +20,12 @@ def add_pymgrid_yaml_representers():
         BatteryDischargeShaper
     )
 
+    from pymgrid.modules.battery.transition_models import (
+        BatteryTransitionModel,
+        BiasedTransitionModel,
+        DecayTransitionModel
+    )
+
 
 def dump_data(data_dict, stream, yaml_tag):
     if not hasattr(stream, "name"):