fix bugs

zhaoziheng · Sep 5, 2024 · ab2792b · ab2792b
1 parent 21b9075
commit ab2792b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/inference.py b/inference.py
@@ -35,7 +35,7 @@ def set_seed(config):
 def main(args):
     # set gpu
     if args.gpu:
-        os.environ['CUDA_VISIBLE_DEVICES'] = args.pin_memory
+        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
 
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
     device=torch.device("cuda", int(os.environ["LOCAL_RANK"]))