Solved by ptrblck in post #2 
              
              
                You could use torch.cuda.device_count() to get the number of GPUs and then torch.cuda.get_device_properties(idx) to get the information of the device using idx.