return it

Author: Cyrilvallez

src/transformers/core_model_loading.py

@@ -609,7 +609,7 @@ def convert_and_load_state_dict_in_model(
 
     model.inverse_converters = inverse_converters
     thread_pool.shutdown(wait=False)
-    return missing_keys, unexpected_keys, mismatch_keys, misc
+    return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, misc
 
 
 # TODO this is not done yet!

src/transformers/modeling_utils.py

@@ -4257,18 +4257,20 @@ def _load_pretrained_model(
             else:
                 raise ValueError("Neither a state dict nor checkpoint files were found.")
 
-            missing_keys, unexpected_keys, mismatched_keys, misc = convert_and_load_state_dict_in_model(
-                model,
-                merged_state_dict,
-                weight_mapping,
-                tp_plan,
-                hf_quantizer,
-                dtype,
-                device_map,
-                model.dtype_plan,
-                device_mesh,
-                disk_offload_index,
-                disk_offload_folder,
+            missing_keys, unexpected_keys, mismatched_keys, disk_offload_index, misc = (
+                convert_and_load_state_dict_in_model(
+                    model,
+                    merged_state_dict,
+                    weight_mapping,
+                    tp_plan,
+                    hf_quantizer,
+                    dtype,
+                    device_map,
+                    model.dtype_plan,
+                    device_mesh,
+                    disk_offload_index,
+                    disk_offload_folder,
+                )
             )
 
             # finally close all opened file pointers