From a72d152b0c8788b954004ee8c01e957701514e32 Mon Sep 17 00:00:00 2001 From: Bratzmeister Date: Tue, 12 Nov 2024 11:53:36 +0000 Subject: [PATCH] fix --cuda-device arg for AMD/HIP devices (#5586) * fix --cuda-device arg for AMD/HIP devices CUDA_VISIBLE_DEVICES is ignored for HIP devices/backend. Instead it uses HIP_VISIBLE_DEVICES. Setting this environment variable has no side effect for CUDA/NVIDIA so it can safely be set in any case and vice versa. * deleted accidental if --- main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/main.py b/main.py index c2321086..05eb31c7 100644 --- a/main.py +++ b/main.py @@ -71,6 +71,7 @@ if os.name == "nt": if __name__ == "__main__": if args.cuda_device is not None: os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) + os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) logging.info("Set cuda device to: {}".format(args.cuda_device)) if args.deterministic: