Skip to content
Snippets Groups Projects
Unverified Commit 2e99ea38 authored by Mashiro's avatar Mashiro Committed by GitHub
Browse files

[Fix]: Fix _save_to_state_dict (#542)

parent c64243aa
No related branches found
No related tags found
No related merge requests found
......@@ -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()
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment