Fix cpu_offload_with_hook code snippet (#1047)

* Fix cpu_offload_with_hook code snippet * Make model explicit for clarity.
huggingface · Feb 8, 2023 · 60460ae · 60460ae
1 parent 978dfc3
commit 60460ae
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
@@ -214,9 +214,9 @@ def cpu_offload_with_hook(
     Example:
 
     ```py
-    hook_1 = cpu_offload_with_hook(model_1, cuda_device)
-    hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
-    hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2)
+    model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
+    model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
+    model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2)
 
     hid_1 = model_1(input)
     for i in range(50):