From bb4feb293ece5d521bfb15c3e9104a2e210ace29 Mon Sep 17 00:00:00 2001
From: Giuseppe Franco <giuseppefranco4@gmail.com>
Date: Mon, 18 Mar 2024 10:16:44 +0100
Subject: [PATCH] Fix (core/zero_point): add return in state_dict (#910)

---
 src/brevitas/core/zero_point.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/brevitas/core/zero_point.py b/src/brevitas/core/zero_point.py
index 58eb73820..872435ec7 100644
--- a/src/brevitas/core/zero_point.py
+++ b/src/brevitas/core/zero_point.py
@@ -283,6 +283,7 @@ def state_dict(self, destination=None, prefix='', keep_vars=False):
         # Avoid saving the init value
         if not self.init_done:
             del output_dict[prefix + 'value']
+        return output_dict
 
     def _load_from_state_dict(
             self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,