fix dilocoo

src/zeroband/checkpoint.py

@@ -223,7 +223,7 @@ def load(self, resume_ckpt_path: str) -> None:
         dcp.load(self.states, checkpoint_id=resume_ckpt_path)
         # since we don't load the param list from the state dict as its the same as the model one we just copy
         if self.diloco_offloaded_param_list is not None:
-            for param_offloaded, param_model in zip(self.diloco_offloaded_param_list, self.model.model.parameters()):
+            for param_offloaded, param_model in zip(self.diloco_offloaded_param_list, self.model.parameters()):
                 param_offloaded.data.copy_(param_model.data)
 
         ## the next part is a fix so that each rank save a different dataloader rank. It not efficient because it reads the state two times from disk