Assign CUDA_VISIBLE_DEVICES if device specified

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"