[Fix]: Fix _save_to_state_dict (#542)

open-mmlab · Sep 22, 2022 · 2e99ea3 · 2e99ea3
1 parent c64243a
commit 2e99ea3
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mmengine/runner/checkpoint.py b/mmengine/runner/checkpoint.py
@@ -610,8 +610,7 @@ def _save_to_state_dict(module, destination, prefix, keep_vars):
         if param is not None:
             destination[prefix + name] = param if keep_vars else param.detach()
     for name, buf in module._buffers.items():
-        # remove check of _non_persistent_buffers_set to allow nn.BatchNorm2d
-        if buf is not None:
+        if buf is not None and name not in module._non_persistent_buffers_set:
             destination[prefix + name] = buf if keep_vars else buf.detach()