From aeed9bd8f74cac4b7aeadc3906f20fbe042d104d Mon Sep 17 00:00:00 2001
From: vfbd <ve.forbryderne@gmail.com>
Date: Fri, 1 Jul 2022 20:16:39 -0400
Subject: [PATCH] Fix base fairseq dense models when using accelerate with a
 GPU

---
 aiserver.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/aiserver.py b/aiserver.py
index 94aa3b60..4346e988 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -642,6 +642,11 @@ def move_model_to_devices(model):
     import breakmodel
 
     if(utils.HAS_ACCELERATE):
+        import accelerate.utils
+        for key, value in model.state_dict().items():
+            target_dtype = torch.float32 if breakmodel.primary_device == "cpu" else torch.float16
+            if(value.dtype is not target_dtype):
+                accelerate.utils.set_module_tensor_to_device(model, key, target_dtype)
         disk_blocks = breakmodel.disk_blocks
         gpu_blocks = breakmodel.gpu_blocks
         ram_blocks = len(utils.layers_module_names) - sum(gpu_blocks)