Hi thanks for the reply, is it possible to set it with python instead of in shell?

Yes, but you need to make sure it is set before initializing the CUDA context. See the code below:
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1" 
torch.cuda.device_count()  # print 1
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  
torch.cuda.device_count()  # still print 1