add .cpu() when gathering tensor

Author: bursteratom

src/transformers/modeling_utils.py

@@ -2944,9 +2944,9 @@ def save_pretrained(
         gathered_state_dict = {}
         for key, value in state_dict.items():
             if hasattr(value, "_local_tensor"):
-                gathered_state_dict[key] = value.to_local()
+                gathered_state_dict[key] = value.to_local().cpu()
             else:
-                gathered_state_dict[key] = value
+                gathered_state_dict[key] = value.cpu()
 
         del state_dict
         state_dict = gathered_state_dict