Fix (stats): add return statement in state_dict

Xilinx · Jan 8, 2024 · 84812f4 · 84812f4
1 parent ade1036
commit 84812f4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/brevitas/core/scaling/standalone.py b/src/brevitas/core/scaling/standalone.py
@@ -223,6 +223,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,