[fix] fix test_checkpoint_io/test_hybrid parallel plugin checkpoint i…

…o.py
duanjunwen · Jun 13, 2024 · cda400b · cda400b
1 parent 0fa84a7
commit cda400b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/colossalai/pipeline/p2p.py b/colossalai/pipeline/p2p.py
@@ -266,7 +266,9 @@ def _send_recv_serialization_object(
     send_object_tensor = None
     if object is not None and send_dst is not None:
         if Version(torch.__version__) >= Version("1.13.0"):
-            send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(object, device=current_device)
+            send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(
+                object, device=current_device, group=c10d._get_default_group()
+            )
         else:
             send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(object)