From 2e99ea3821d4a99bbecd40865c38c57696fc1408 Mon Sep 17 00:00:00 2001
From: Mashiro <57566630+HAOCHENYE@users.noreply.github.com>
Date: Thu, 22 Sep 2022 14:10:35 +0800
Subject: [PATCH] [Fix]: Fix _save_to_state_dict (#542)
---
mmengine/runner/checkpoint.py | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mmengine/runner/checkpoint.py b/mmengine/runner/checkpoint.py
index 44ae462d..9b80c4e0 100644
--- 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()
--
GitLab