Skip to content

GitLab

Commit 4c54fd2b authored by Thor Johnsen's avatar Thor Johnsen
Browse files

Make step global state variable

parent 5caf95ca
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
apex/contrib/optimizers/distributed_fused_lamb.py
View file @ 4c54fd2b
...@@ -72,7 +72,8 @@ class DistributedFusedLAMB(torch.optim.Optimizer): ...@@ -72,7 +72,8 @@ class DistributedFusedLAMB(torch.optim.Optimizer):
dwu_group_size=0, dwu_num_blocks=4, dwu_num_chunks=4, dwu_group_size=0, dwu_num_blocks=4, dwu_num_chunks=4,
dwu_num_rs_pg=1, dwu_num_ar_pg=4, dwu_num_ag_pg=0, dwu_num_rs_pg=1, dwu_num_ar_pg=4, dwu_num_ag_pg=0,
e5m2_allgather=False): e5m2_allgather=False):
global distributed_lamb_cuda global fused_adam_cuda, distributed_lamb_cuda
fused_adam_cuda = importlib.import_module("fused_adam_cuda")
distributed_lamb_cuda = importlib.import_module("distributed_lamb_cuda") distributed_lamb_cuda = importlib.import_module("distributed_lamb_cuda")
self._amp_scale_adjustment = amp_scale_adjustment self._amp_scale_adjustment = amp_scale_adjustment
......
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