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

fix the scatter when input is cpu #1621

Merged
merged 5 commits into from
Jan 24, 2022
Merged

fix the scatter when input is cpu #1621

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e593bc7
fix the scatter when input is cpu
Bilibilee Dec 27, 2021
5ef7870
Update _functions.py
Bilibilee Jan 17, 2022
903fb1f
add unittests
zhouzaida Jan 22, 2022
18cc4eb
add a blank line
zhouzaida Jan 22, 2022
ff776e1
fix unittet
zhouzaida Jan 22, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions mmcv/parallel/_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,7 @@ def scatter(input, devices, streams=None):
if devices != [-1]:
with torch.cuda.device(devices[0]), torch.cuda.stream(stream):
output = output.cuda(devices[0], non_blocking=True)
else:
# unsqueeze the first dimension thus the tensor's shape is the
# same as those scattered with GPU.
output = output.unsqueeze(0)

return output
else:
raise Exception(f'Unknown type {type(input)}.')
Expand Down Expand Up @@ -76,4 +73,4 @@ def forward(target_gpus, input):
if streams is not None:
synchronize_stream(outputs, target_gpus, streams)

return tuple(outputs)
return tuple(outputs) if isinstance(outputs, list) else (outputs, )