remove the to in the copy because jack said so and the test pass

PrimeIntellect-ai · Sep 27, 2024 · 9d801d3 · 9d801d3
1 parent b43bde7
commit 9d801d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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]:
         """