Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[NFC] polish colossalai/nn/parallel/data_parallel.py code style #1570

Merged
merged 1 commit into from
Sep 8, 2022
Merged

[NFC] polish colossalai/nn/parallel/data_parallel.py code style #1570

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 2 additions & 2 deletions colossalai/nn/parallel/data_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,7 @@ def _save_to_state_dict(self, destination, prefix, keep_vars, record_flag: bool
for tensor in chunk.get_tensors():
rec_p = torch.empty([0])
if record_flag:
rec_p = tensor.cpu() # move the whole tensor to CPU mem
rec_p = tensor.cpu() # move the whole tensor to CPU mem
assert tensor not in param_to_save_data
param_to_save_data[tensor] = rec_p
# release the actual memory of the chunk
Expand Down Expand Up @@ -406,7 +406,7 @@ def load_state_dict(self, state_dict: 'OrderedDict[str, torch.Tensor]', strict:
state_dict = state_dict.copy()
if metadata is not None:
# mypy isn't aware that "_metadata" exists in state_dict
state_dict._metadata = metadata # type: ignore[attr-defined]
state_dict._metadata = metadata # type: ignore[attr-defined]

prefix = ''
local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
Expand Down