在TensorFlow中可以这样配置GPU显存，在pytorch中如何实现呢？

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
    config = tf.ConfigProto(allow_soft_placement=False, log_device_placement=False)
    config.gpu_options.per_process_gpu_memory_fraction = 0.6  # 分配百分之七十的显存给程序使用，避免内存溢出，可以自己调整
    config.gpu_options.allow_growth = True  # 按需分配显存，这个比较重要
    with tf.Session(config=config) as sess: