Skip to content

Instantly share code, notes, and snippets.

@rahulvigneswaran

rahulvigneswaran/GPU_selector.py

Created November 8, 2019 10:45
Show Gist options
  • Save rahulvigneswaran/ed2137c6086bdd4bdbb3b3c60384dff5 to your computer and use it in GitHub Desktop.
Save rahulvigneswaran/ed2137c6086bdd4bdbb3b3c60384dff5 to your computer and use it in GitHub Desktop.
Download ZIP
Select GPU in Pytorch
Raw
GPU_selector.py
import os
parser = argparse.ArgumentParser()
parser.add_argument("--gpu", default="0", type=str)
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment