do copy instead of assigment

PrimeIntellect-ai · Sep 27, 2024 · a143cab · a143cab
1 parent ba392de
commit a143cab
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 = param_offloaded.data.to("cuda")  # todo: use copy_ here
+            param.data.copy_(param_offloaded.data.to(param.device))  # todo: use copy_ here
 
     def get_offloaded_param(self, model: nn.Module) -> list[nn.Parameter]:
         """