Skip to content

GitLab

Unverified Commit 1fa778d7 authored by Benjamin Lefaudeux's avatar Benjamin Lefaudeux Committed by GitHub
Browse files

[fix] Lightning compatibility (#510)

parent 5e8a6422
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
fairscale/optim/oss.py
View file @ 1fa778d7
...@@ -400,7 +400,7 @@ class OSS(Optimizer): ...@@ -400,7 +400,7 @@ class OSS(Optimizer):
# Populate the sharded optimizer state on the fly, # Populate the sharded optimizer state on the fly,
# remove the params that this rank does not own # remove the params that this rank does not own
if self.param_to_rank[param] != self.rank: if self.param_to_rank[param] != self.rank:
state_dict["state"][key] = None state_dict["state"][key] = {}
else: else:
self.optim.state[param] = recursive_copy_to_device(value, non_blocking=True, device=param.device) self.optim.state[param] = recursive_copy_to_device(value, non_blocking=True, device=param.device)
else: else:
......
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