Revert "Make sure #126704 is BC for torch.save-ed `nn.Module` (#128344)"
This reverts commit 841d87177a900c2bbd59b6589165189141c4e8bb.
Reverted https://github.com/pytorch/pytorch/pull/128344 on behalf of https://github.com/clee2000 due to broke internal typecheck D58394110 (which probably means the code wouldn't work either but I guess it didn't run on the diff). Probably an easy fix? ([comment](https://github.com/pytorch/pytorch/pull/126704#issuecomment-2161299193))
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
index 2c2c286..942d5b8 100644
--- a/torch/nn/modules/module.py
+++ b/torch/nn/modules/module.py
@@ -1955,12 +1955,7 @@
for name, module in self._modules.items():
if module is not None:
module.state_dict(destination=destination, prefix=prefix + name + '.', keep_vars=keep_vars)
- for value in self._state_dict_hooks.values():
- # For BC reasons
- if isinstance(value, tuple):
- hook, from_private = value
- else:
- hook, from_private = value, True
+ for (hook, from_private) in self._state_dict_hooks.values():
hook_result = hook(self, destination, prefix, local_metadata)
if from_private and hook_result is not None:
destination = hook_result
