diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py
index 9f7a2466..cdf30a0d 100644
--- a/src/twinkle/checkpoint_engine/manager.py
+++ b/src/twinkle/checkpoint_engine/manager.py
@@ -127,8 +127,6 @@ def decide_backend_engine(
 
         platform_name = Platform.get_platform(platform).__name__
         if mode == 'colocate':
-            if platform_name != 'GPU':
-                raise NotImplementedError("mode='colocate' currently requires the GPU platform.")
             from twinkle.checkpoint_engine import IPCCheckpointEngine
             return IPCCheckpointEngine
         if mode != 'standalone':
@@ -166,8 +164,8 @@ def sync_weights(self, merge_and_sync=True):
             self._sync_weights_naive(merge_and_sync)
             return
 
-        is_master = [True] + [False] * (self.model.device_mesh.world_size - 1)
-        model_metadata = self.model.prepare_checkpoint_engine(is_master)
+        model_metadata = self.model.prepare_checkpoint_engine([True]
+                                                              + [False] * (self.model.device_mesh.world_size - 1))
         self.sampler.prepare_checkpoint_engine(False)
         model_kwargs, sampler_kwargs = self.backend_cls.build_topology(
             self.model.device_mesh.world_size,