MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / state_dict_hook

Method state_dict_hook

python/oneflow/test/modules/test_module.py:436–444  ·  view source on GitHub ↗
(module, destination, prefix, local_metadata)

Source from the content-addressed store, hash-verified

434 destination_check = None
435
436 def state_dict_hook(module, destination, prefix, local_metadata):
437 for submodule_name, submodule in module.named_modules():
438 for attr_name, attr in submodule.__dict__.items():
439 if isinstance(attr, torch.Tensor):
440 mod_prefix = prefix + submodule_name
441 key = mod_prefix + ("." if mod_prefix else "") + attr_name
442 destination[key] = attr
443 nonlocal destination_check
444 destination_check = destination
445
446 class CustomModule(flow.nn.Module):
447 def __init__(self):

Callers

nothing calls this directly

Calls 2

named_modulesMethod · 0.80
itemsMethod · 0.45

Tested by

no test coverage detected