class_docstr: Base test class and used for CPU case.
docstr: Test that original (unwrapped) state_dict can be loaded in wrapped module.
docstr_tok: ["Test","that","original","(","unwrapped",")","state_dict","can","be","loaded","in","wrapped","module","."]
code: def test_load_state_dict(self):
      """Test that original (unwrapped) state_dict can be loaded in wrapped module."""
      for module_init_fn in self._get_module_init_fns():
          module = module_init_fn()
          ref_state_dict = module.state_dict()
          ref_output = self._compute_output(module)
  
          module = module_init_fn(seed=1234)
          flat_module = FlattenParamsWrapper(
              module,
              param_list=module.param_list,
              flat_param_names=module.flat_param_names,
          )
  
          # This should work without the unflatten_params context manager
          flat_module.load_state_dict(ref_state_dict)
          flat_output = self._compute_output(flat_module)
          assert objects_are_equal(ref_output, flat_output)
  
          # And it should work with the context manager too
          with flat_module.unflatten_params():
              flat_module.load_state_dict(ref_state_dict)
          flat_output = self._compute_output(flat_module)
          assert objects_are_equal(ref_output, flat_output)
  
imports: from collections import OrderedDict
  import unittest
  import torch
  from fairscale.fair_dev.testing.testing import objects_are_equal
  from fairscale.nn import FlattenParamsWrapper
label_desc: Use image node_id 22 for calling the TestFlattenParams obj's underlying member method code with example usage: obj.test_load_state_dict() without return types