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

Method test_state_dict

python/oneflow/test/modules/test_module.py:217–232  ·  view source on GitHub ↗
(test_case)

Source from the content-addressed store, hash-verified

215
216 @flow.unittest.skip_unless_1n1d()
217 def test_state_dict(test_case):
218 class CustomModule(flow.nn.Module):
219 def __init__(self, param1, param2):
220 super().__init__()
221 self.param1 = param1
222 self.param2 = param2
223
224 tensor0 = flow.nn.Parameter(flow.Tensor(2, 3))
225 tensor1 = flow.nn.Parameter(flow.Tensor(2, 3))
226 sub_module = CustomModule(tensor0, tensor1)
227 m = CustomModule(tensor1, sub_module)
228 state_dict = m.state_dict()
229 test_case.assertEqual(
230 state_dict,
231 {"param2.param1": tensor0, "param2.param2": tensor1, "param1": tensor1},
232 )
233
234 @flow.unittest.skip_unless_1n1d()
235 def test_parameter(test_case):

Callers

nothing calls this directly

Calls 4

ParameterMethod · 0.80
CustomModuleClass · 0.70
TensorMethod · 0.45
state_dictMethod · 0.45

Tested by

no test coverage detected