fix

Author: Cyrilvallez

src/transformers/core_model_loading.py

@@ -394,7 +394,6 @@ def set_param_for_module(
     with log_to_misc(full_param_name, misc, full_param_name):
         module_path, _, param_name = full_param_name.rpartition(".")
         module_obj = model.get_submodule(module_path) if module_path else model
-        param_value = param_value[0] if isinstance(param_value, list) else param_value[...]
         ref = getattr(module_obj, param_name)
 
         use_dtensor = hasattr(distributed_operation, "use_dtensor") and distributed_operation.use_dtensor
@@ -584,6 +583,7 @@ def convert_and_load_state_dict_in_model(
                                     )
 
                         for k, output_value in realized_value.items():
+                            output_value = output_value[0] if isinstance(output_value, list) else output_value
                             for src in converter.source_keys:  # what should happen to k when we meet k at saving
                                 inverse_converters[k] = {src: converter}
 

tests/test_modeling_common.py

@@ -2371,7 +2371,7 @@ def test_disk_offload_bin(self):
                 max_size = int(self.model_split_percents[1] * model_size)
                 max_memory = {0: max_size, "cpu": max_size}
                 new_model = model_class.from_pretrained(
-                    tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir
+                    tmp_dir, device_map="auto", max_memory=max_memory, offload_folder=tmp_dir, use_safetensors=False
                 )
 
                 self.check_device_map_is_respected(new_model, new_model.hf_device_map)