accelerate: [Include example code] mixed_precision="fp16" will break torch.save function.

System Info

accelerate-0.14.0
Python 3.7.15
Pytorch 1.12.1+cu113

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

from accelerate import Accelerator
import torch
import torch.nn as nn

class ExampleModule(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.conv = nn.Conv2d(3, 3, kernel_size=1)

model = ExampleModule()

#mixed_precision="fp16" will give error on torch.save
#mixed_precision="no" will work with torch.save
accelerator = Accelerator(
        gradient_accumulation_steps=1,
        mixed_precision="fp16",
        log_with="tensorboard",
        logging_dir=".",
    )

#Always work
torch.save(model,  "/model_original.model")

#Will break torch.save if the model if mixed_precision="fp16" 
model = accelerator.prepare(model)

#Error with mixed_precision="fp16" 
torch.save(model,  "/model_acc.model")
#Error as well with mixed_precision="fp16" 
torch.save(accelerator.unwrap_model(model),  "/model_unwrap.sd")

It will return this error if mixed_precision=“fp16”

---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
[<ipython-input-1-5ce45839c137>](https://localhost:8080/#) in <module>
     27 
     28 #Error
---> 29 torch.save(model,  "/model_acc.model")
     30 #Error as well
     31 torch.save(accelerator.unwrap_model(model),  "/model_unwrap.sd")

1 frames
[/usr/local/lib/python3.7/dist-packages/torch/serialization.py](https://localhost:8080/#) in _save(obj, zip_file, pickle_module, pickle_protocol)
    587     pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
    588     pickler.persistent_id = persistent_id
--> 589     pickler.dump(obj)
    590     data_value = data_buf.getvalue()
    591     zip_file.write_record('data.pkl', data_value, len(data_value))

PicklingError: Can't pickle <function _forward_unimplemented at 0x7fb39d0b0320>: it's not the same object as torch.nn.modules.module._forward_unimplemented

Expected behavior

torch.save should work even if Accelerator is set to fp16