Skip to content

GitLab

Unverified Commit 20ab1f55 authored by Frank Lee's avatar Frank Lee Committed by GitHub
Browse files

[bug] fixed broken test_found_inf (#725)

parent 193dc8da
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/test_zero_data_parallel/test_found_inf.py
View file @ 20ab1f55
...@@ -31,7 +31,7 @@ def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio) ...@@ -31,7 +31,7 @@ def _run_test_found_inf(cpu_offload, shard_strategy_class, gpu_margin_mem_ratio)
model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func() model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
with ZeroInitContext( with ZeroInitContext(
target_device=torch.device(f'cpu:0') if cpu_offload else torch.device(f'cuda:{get_current_device()}'), target_device=torch.device(f'cpu:0') if cpu_offload else torch.device(get_current_device()),
shard_strategy=shard_strategy, shard_strategy=shard_strategy,
shard_param=True): shard_param=True):
zero_model = model_builder(checkpoint=True) zero_model = model_builder(checkpoint=True)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment