From 2a7513e869e8cd5d8b2e14daf396b2d1cc85abde Mon Sep 17 00:00:00 2001
From: dcarron <daniel.carron@idiap.ch>
Date: Mon, 17 Apr 2023 16:54:31 +0200
Subject: [PATCH] Assign CUDA_VISIBLE_DEVICES if device specified
---
src/ptbench/engine/trainer.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/src/ptbench/engine/trainer.py b/src/ptbench/engine/trainer.py
index 896ce9a3..723c6d66 100644
--- a/src/ptbench/engine/trainer.py
+++ b/src/ptbench/engine/trainer.py
@@ -42,6 +42,7 @@ class AcceleratorProcessor:
if len(split_accelerator) > 1:
devices = split_accelerator[1:]
devices = [int(i) for i in devices]
+ os.environ["CUDA_VISIBLE_DEVICES"] = devices
else:
devices = "auto"
--
GitLab