From 9d801d31829def5a7e8fe23668a566e6e7d79cb4 Mon Sep 17 00:00:00 2001
From: Sami Jaghouar <sami.jaghouar@hotmail.fr>
Date: Fri, 27 Sep 2024 04:38:57 +0000
Subject: [PATCH] remove the to in the copy because jack said so and the test
 pass

---
 src/zeroband/diloco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/zeroband/diloco.py b/src/zeroband/diloco.py
index e7113918..b132c707 100644
--- a/src/zeroband/diloco.py
+++ b/src/zeroband/diloco.py
@@ -107,7 +107,7 @@ def sync_inner_model(self, model: nn.Module):
 
         self._logger.debug("sync inner model")
         for param_offloaded, param in zip(self.param_list_cpu, model.parameters()):
-            param.data.copy_(param_offloaded.data.to(param.device))  # todo: use copy_ here
+            param.data.copy_(param_offloaded.data)  # todo: use copy_ here
 
     def get_offloaded_param(self, model: nn.Module) -> list[nn.Parameter]:
         """