Automatically move data to the GPU?

1

We’ve got the follow block of code repeated a few places in our project’s code base. I’d like to refactor it so that so much code isn’t repeated.

        for i_batch, (sample_batch, local_l_mean, local_l_var, batch_index, _) in enumerate(our_data_loader):
            if using_cuda:
                sample_batch = sample_batch.cuda(async=True)
                local_l_mean = local_l_mean.cuda(async=True)
                local_l_var = local_l_var.cuda(async=True)
                batch_index = batch_index.cuda(async=True)

Is there a flag we can pass to DataLoader so that all the tensors it yields are automatically copied to the GPU? (i.e., without having to call .cuda() for each tensor.) Or if not, any suggestions for the best way to avoid repeating this block of code everywhere that we iterate over a dataset?

EDIT: I suppose I could call function like the following each time within the loop:

def to_cuda(tensor_list):
     return [t.cuda(async=True) for t in tensor_list]

Please let me know if there’s a better way.

Thanks!