Skip to content

Commit

Permalink
SA: for #958: removing root gpu hack in trainer/evaluation_loop
Browse files Browse the repository at this point in the history
  • Loading branch information
shubhamagarwal92 authored and williamFalcon committed Apr 3, 2020
1 parent 5fc8165 commit 7fa8471
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion pytorch_lightning/trainer/evaluation_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -429,9 +429,13 @@ 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]
else:
raise RuntimeError(
'Expected `data_parallel_device_ids` as a list, cannot determine root gpu.'
)
batch = self.transfer_batch_to_gpu(batch, root_gpu)
args[0] = batch

Expand Down

0 comments on commit 7fa8471

Please sign in to comment.