From f2e113fec3807c8bee13ed96f371a7d513acbfd5 Mon Sep 17 00:00:00 2001
From: ahalev <avishaihalev@gmail.com>
Date: Tue, 4 Apr 2023 22:56:39 -0700
Subject: [PATCH] add battery transition models to representers

---
 src/pymgrid/utils/serialize.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/pymgrid/utils/serialize.py b/src/pymgrid/utils/serialize.py
index 45be9c2c..69d8a308 100644
--- 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"):