diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d6a323559..d614b47e2 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
             if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
                 model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
             elif self.mixed_precision == "bf16" and self.distributed_type != DistributedType.TPU:
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 7ba4482cc..3ebea2eef 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -17,6 +17,7 @@
 """
 
 
+import pickle
 from functools import update_wrapper
 from typing import Any, Mapping
 
@@ -473,8 +474,6 @@ class ConvertOutputsToFp32:
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
     Args:
         model_forward (`Callable`):
             The function which outputs we want to treat.
@@ -490,6 +489,11 @@ def __init__(self, model_forward):
     def __call__(self, *args, **kwargs):
         return convert_to_fp32(self.model_forward(*args, **kwargs))
 
+    def __getstate__(self):
+        raise pickle.PicklingError(
+            "Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it."
+        )
+
 
 convert_outputs_to_fp32 = ConvertOutputsToFp32
 
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index a4a7b6312..1eb7cca88 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -21,7 +21,6 @@
 from ..state import AcceleratorState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
-from .operations import ConvertOutputsToFp32
 
 
 if is_deepspeed_available():
@@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break
+            model.forward = forward
     return model
 
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 1e9d18c17..7c7629a1a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -19,8 +19,15 @@
 
 import torch
 
+from accelerate.test_utils.testing import require_cuda
 from accelerate.test_utils.training import RegressionModel
-from accelerate.utils import convert_outputs_to_fp32, find_device, patch_environment, send_to_device
+from accelerate.utils import (
+    convert_outputs_to_fp32,
+    extract_model_from_parallel,
+    find_device,
+    patch_environment,
+    send_to_device,
+)
 
 
 ExampleNamedTuple = namedtuple("ExampleNamedTuple", "a b c")
@@ -74,9 +81,20 @@ def test_patch_environment(self):
         self.assertNotIn("AA", os.environ)
         self.assertNotIn("BB", os.environ)
 
-    def test_convert_to_32_lets_model_pickle(self):
+    def test_can_undo_convert_outputs(self):
         model = RegressionModel()
+        model._original_forward = model.forward
         model.forward = convert_outputs_to_fp32(model.forward)
+        model = extract_model_from_parallel(model)
+        _ = pickle.dumps(model)
+
+    @require_cuda
+    def test_can_undo_fp16_conversion(self):
+        model = RegressionModel()
+        model._original_forward = model.forward
+        model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)
+        model.forward = convert_outputs_to_fp32(model.forward)
+        model = extract_model_from_parallel(model)
         _ = pickle.dumps(model)
 
     def test_find_device(self):


Review comments: body: @sgugger I know how you feel about `_` but I felt this was a decent time for one because the user really *shouldn't* need to know this exists at all or try to accidentally call it in any way
diff_hunk: @@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
body: Agree in this case :-)
diff_hunk: @@ -926,6 +926,7 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
             model = torch.nn.parallel.DistributedDataParallel(model, **kwargs)
         if self.native_amp:
+            model._original_forward = model.forward
body: Are we sure this does not make a regression with this issue?
diff_hunk: @@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
body: The issue will always exist because it can never be pickled if we've wrapped it in a mixed precision decorator. So it will always have to go through `unwrap_model`. And this only ever comes up if we do mixed precision. See the nb shown in the bug report here: https://colab.research.google.com/drive/11fvrk1Jslw2VIRTF6h0pdGgJx5AkAIMv?usp=sharing

Where specifically our explicit wrapping around `torch.cuda.amp` causes issues, and technically was the original issue too.

Or, phrasing it another way, the original issue doesn't come from Accelerate actually it comes from torch and using the autocast decorator. To keep things consistent this just makes them all use decorators so that we can quickly just unwrap all of them. 
diff_hunk: @@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
body: I can try to see if it's possible to inject a warning before trying to pickle/save it and tell the user they need to unwrap the model first. 
diff_hunk: @@ -468,13 +468,11 @@ def _is_fp16_bf16_tensor(tensor):
     return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_bf16_tensor)
 
 
-class ConvertOutputsToFp32:
+def convert_outputs_to_fp32(model_forward):
     """
     Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
     precision will be convert back to FP32.
 
-    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).
-
body: Let's put your comment on the PR here as well, to explain what this code is doing 🙏 
diff_hunk: @@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break
body: A custom `__getstate__` func let's us raise an error when the user tries to prepare a model that had automatic mixed precision performed, ensuring with a clear error what the proper behavior needs to be, and the user doesn't have to know that this is due to PyTorch *and* accelerate and their behaviors
diff_hunk: @@ -490,6 +488,11 @@ def __init__(self, model_forward):
     def __call__(self, *args, **kwargs):
         return convert_to_fp32(self.model_forward(*args, **kwargs))
 
+    def __getstate__(self):
+        raise pickle.PicklingError(
+            "Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it."
+        )
+
body: When a function is wrapped using functools (such that we have `torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)` the wrapped function (the original `model.forward`) exists in the `__wrapped__` attribute. When using `functools.update_wrapper` (such as in `ConvertOutputsToFp32`) this will also in turn add a new `__wrapped__` function, allowing us to follow the chain of `__wrapped__` until we return to the original model forward function, which was stored away earlier.
diff_hunk: @@ -53,8 +52,13 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = False):
 
     if not keep_fp32_wrapper:
         forward = getattr(model, "forward")
-        if isinstance(forward, ConvertOutputsToFp32):
-            setattr(model, "forward", forward.model_forward)
+        original_forward = model.__dict__.pop("_original_forward", None)
+        if original_forward is not None:
+            while hasattr(forward, "__wrapped__"):
+                forward = forward.__wrapped__
+                if forward == original_forward:
+                    break

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: Some notes on the final version (which will also be trickled into the code of the PR):

- I kept this as a class so that we can implement a custom `__getstate__` which will return an explicit pickling error if someone tries to pickle a wrapped model with automatic mixed precision on it
- When a function is wrapped using functools (such that we have `torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)` the wrapped function (the original `model.forward`) exists in the `__wrapped__` attribute. When using `functools.update_wrapper` (such as in `ConvertOutputsToFp32`) this will also in turn add a new `__wrapped__` function, allowing us to follow the chain of `__wrapped__` until we return to the original model forward function, which was stored away earlier.