class_docstr: Base test class and used for CPU case.
docstr: Test that flat state dict can be reloaded and produces the same results.
docstr_tok: ["Test","that","flat","state","dict","can","be","reloaded","and","produces","the","same","results","."]
code: def test_flat_state_dict(self):
      """Test that flat state dict can be reloaded and produces the same results."""
      for module_init_fn in self._get_module_init_fns():
          orig_module = module_init_fn()
          flat_module = FlattenParamsWrapper(
              orig_module,
              param_list=orig_module.param_list,
              flat_param_names=orig_module.flat_param_names,
          )
          ref_output = self._compute_output(flat_module)
  
          flat_state_dict = flat_module.flat_state_dict()
  
          orig_module = module_init_fn(seed=1234)
          new_module = FlattenParamsWrapper(
              orig_module,
              param_list=orig_module.param_list,
              flat_param_names=orig_module.flat_param_names,
          )
          new_module.load_state_dict(flat_state_dict)
          new_output = self._compute_output(new_module)
  
          assert objects_are_equal(ref_output, new_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 23 for calling the TestFlattenParams obj's underlying member method code with example usage: obj.test_flat_state_dict() without return types