-
Notifications
You must be signed in to change notification settings - Fork 3.5k
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.
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
setting PGU device #1128
setting PGU device #1128
Changes from all commits
5c554a1
2f17b2f
6d89505
54e9a5e
83c291d
8a66308
ae30b27
67072ad
c684cc9
3b42a40
30f0bf3
9473303
10a5b05
dd8a7c2
2acd85d
08bb574
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -421,9 +421,18 @@ def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test_mode: | |
# single GPU data transfer | ||
if self.single_gpu: | ||
# for single GPU put inputs on gpu manually | ||
root_gpu = 0 | ||
|
||
if isinstance(self.data_parallel_device_ids, list): | ||
root_gpu = self.data_parallel_device_ids[0] | ||
|
||
# set cuda device to root gpu | ||
root_device = (torch.device("cuda", root_gpu) | ||
if root_gpu >= 0 else torch.device("cpu")) | ||
torch.cuda.set_device(root_device) | ||
else: | ||
raise RuntimeError( | ||
'Expected `data_parallel_device_ids` as a list, cannot determine root gpu.' | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Above we have There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was already an if clause: |
||
) | ||
batch = self.transfer_batch_to_gpu(batch, root_gpu) | ||
args[0] = batch | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
To me it looks like this should not go here, because the method is called "determine_root_gpu_device", and the added code also checks if the device is cpu. This code should probably go to the place where we call
determine_root_gpu_device
.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Where do you suggest to place this?