mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix]: Fix _save_to_state_dict (#542)

Browse Source
This commit is contained in:
Mashiro 2022-09-22 14:10:35 +08:00 committed by GitHub
parent c64243aa9e
commit 2e99ea3821
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mmengine/runner/checkpoint.py
View File

@ -610,8 +610,7 @@ def _save_to_state_dict(module, destination, prefix, keep_vars):
if param is not None: if param is not None:
destination[prefix + name] = param if keep_vars else param.detach() destination[prefix + name] = param if keep_vars else param.detach()
for name, buf in module._buffers.items(): for name, buf in module._buffers.items():
# remove check of _non_persistent_buffers_set to allow nn.BatchNorm2d if buf is not None and name not in module._non_persistent_buffers_set:
if buf is not None:
destination[prefix + name] = buf if keep_vars else buf.detach() destination[prefix + name] = buf if keep_vars else buf.detach()