fix: fix broadcast_pyobj breaking VerlEngine (#5997)
This commit is contained in:
@@ -897,7 +897,10 @@ def broadcast_pyobj(
|
||||
src: int = 0,
|
||||
force_cpu_device: bool = True,
|
||||
):
|
||||
"""Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
|
||||
"""Broadcast inputs from src rank to all other ranks with torch.dist backend.
|
||||
The `rank` here refer to the source rank on global process group (regardless
|
||||
of dist_group argument).
|
||||
"""
|
||||
device = torch.device(
|
||||
"cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user