Skip to content

GitLab

Commit 2c744ee5 authored by Thor Johnsen's avatar Thor Johnsen
Browse files

Bug(?) fix

parent 208c91e0
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
apex/contrib/optimizers/distributed_fused_adam.py
View file @ 2c744ee5
...@@ -256,7 +256,7 @@ class DistributedFusedAdam(torch.optim.Optimizer): ...@@ -256,7 +256,7 @@ class DistributedFusedAdam(torch.optim.Optimizer):
def _pipeline_block_step(self, block_id): def _pipeline_block_step(self, block_id):
if self._new_params is None: if self._new_params is None:
self._new_params = torch.zeros_like(self._flat_grads,dtype=uint8 if self._e5m2_allgather else self._flat_grads.dtype) self._new_params = torch.zeros_like(self._flat_grads,dtype=torch.uint8 if self._e5m2_allgather else self._flat_grads.dtype)
start = block_id * self._block_size start = block_id * self._block_size
end = start + self._block_size end = start + self._block_size
......
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