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(): flat_module = FlattenParamsWrapper(module_init_fn()) ref_output = self._get_output(flat_module) flat_state_dict = flat_module.flat_state_dict() new_module = FlattenParamsWrapper(module_init_fn(seed=1234)) new_module.load_state_dict(flat_state_dict) new_output = self._get_output(new_module) assert objects_are_equal(ref_output, new_output)
def test_flat_state_dict(self): flat_module = self._get_shared_params_transformer() flat_module = FlattenParamsWrapper(flat_module) ref_output = self._get_output(flat_module) flat_state_dict = flat_module.flat_state_dict() new_module = self._get_shared_params_transformer(seed=1234) new_module = FlattenParamsWrapper(new_module) new_module.load_state_dict(flat_state_dict) new_output = self._get_output(new_module) assert objects_are_equal(ref_output, new_output)