from absl import app
from absl import flags

FLAGS = flags.FLAGS

flags.DEFINE_string('gpu', None, 'comma separated list of GPU(s) to use.')

def main(argv):
del argv
if FLAGS.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
else:
print('Please assign GPUs.')
exit()

if __name__ == '__main__':
app.run(main)

 

python test.py --gpu 0

黄世宇/Shiyu Huang's Personal Page:​​https://huangshiyu13.github.io/​