Skip to content

Commit

Permalink
Specifiy the devices when registering the backend to avoid warnings (#16
Browse files Browse the repository at this point in the history
)

pytest process_group_test.py
  • Loading branch information
fegin authored Nov 23, 2024
1 parent 829d26c commit c7a7f1e
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion torchft/process_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,11 @@ def create_pg(
) -> ProcessGroup:
return self

dist.Backend.register_backend(group_name, create_pg)
if torch.cuda.is_available():
devices = ["cuda", "cpu"]
else:
devices = ["cpu"]
dist.Backend.register_backend(group_name, create_pg, devices=devices)

return dist.new_group(
ranks=[dist.get_rank()],
Expand Down

0 comments on commit c7a7f1e

Please sign in to comment.