okay, another doubt, after I stored name of gpu’s in the device
device = torch.cuda.get_device_name(range(device_id))  then how to use it in model like in the following calling -
if torch.cuda.device_count()>1:
   model = nn.DataParallel(model)
   model = model.to(device)
elif train_on_gpu:
   model = model.to(device)