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,