okay, so I am trying to run my model on 2 gpu’s I have. Following is my code and it is giving me
Type Error: '<' not supported between instances of 'range' and 'int' 
device_id = torch.cuda.device_count()
device = torch.cuda.get_device_name(range(device_id))

if torch.cuda.device_count()>1:
   model = nn.DataParallel(model)
   model = model.to(device)
elif train_on_gpu:
   model = model.to(device)