@@ -163,4 +172,2 @@
         for name in ms_params_name:
-            _param_list = []
             ms_params = getattr(self, name)
-            for p in ms_params:

@@ -173,0 +183,4 @@
+                    _load(ms_params[i], state_tensor[i])
+            else:
+                _data = cast_to_ms_tensor(state_tensor)
+                ms_params.set_data(_data)

@@ -163,0 +167,4 @@
+                for p in ms_params:
+                    _state.append(_save(p))
+            else:
+                _state = tensor(ms_params.asnumpy())

@@ -173,0 +182,4 @@
+            if isinstance(ms_params, Iterable):
+                if not isinstance(state_tensor, Iterable) or len(ms_params) != len(state_tensor):
+                    raise ValueError(f"state_dict of ms_param '{name}' is not correct. please check. "
+                                     f"(ms_param '{name}' is Iterable, but state_dict['{name}'] is not.)")