How to set the right GPU

vision
1

How to set the GPU at iter? I want that iter initializes on cuda:0 on the second line.

device = "cuda:0"
train_adv_loader_iter = iter(train_adv_loader) # How to set the GPU here? 
for i, (images, target) in enumerate( tqdm(train_loader)):
    data_adv = next(train_adv_loader_iter)

    print( data_adv.get_device() ) # cuda:2
    images_adv = data_adv[0].to(device)
    target_adv   = data_adv[1].to(device)
2

I set the visible device on the very top of the of the code.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

other way could be you can pass in the device information into your data loader so that the correct device is used in that.

3

Thanks. This is a very easy solution.