Skip to content

Commit

Permalink
fix tests..
Browse files Browse the repository at this point in the history
  • Loading branch information
misko committed Jul 24, 2024
1 parent 8508e2f commit 4426870
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions tests/core/common/test_gp_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
)
from fairchem.core.common.test_utils import (
PGConfig,
_init_pg_and_rank_and_launch_test,
init_pg_and_rank_and_launch_test,
spawn_multi_process,
)

Expand All @@ -26,7 +26,7 @@ def test_basic_setup(world_size: int, input: torch.Tensor, expected_output: list
backend="gloo", world_size=world_size, gp_group_size=1, use_gp=True
)
output = spawn_multi_process(
config, _dummy_call, _init_pg_and_rank_and_launch_test, input
config, _dummy_call, init_pg_and_rank_and_launch_test, input
)
assert output == expected_output

Expand Down Expand Up @@ -64,7 +64,7 @@ def test_scatter_tensors(
output = spawn_multi_process(
config,
scatter_to_model_parallel_region,
_init_pg_and_rank_and_launch_test,
init_pg_and_rank_and_launch_test,
input,
)
for out, expected_out in zip(output, expected_output):
Expand Down Expand Up @@ -112,7 +112,7 @@ def test_gather_tensors(
backend="gloo", world_size=world_size, gp_group_size=gp_size, use_gp=True
)
output = spawn_multi_process(
config, scatter_gather_fn, _init_pg_and_rank_and_launch_test, input
config, scatter_gather_fn, init_pg_and_rank_and_launch_test, input
)
for out, expected_out in zip(output, expected_output):
assert torch.equal(out, expected_out)

0 comments on commit 4426870

Please sign in to comment.