hip_visible_devices

2026-09-19 17:24:32 +02:00 · 2024-08-13 15:03:25 +09:00
parent 26d1d42b06
commit c6cd0723cf
1 changed files with 2 additions and 0 deletions
@@ -502,6 +502,8 @@ def install_rocm_zluda(torch_command):
                error = e
                log.warning(f'Failed to install ZLUDA: {e}')
            if error is None:
+                if args.device_id is not None:
+                    os.environ['HIP_VISIBLE_DEVICES'] = args.device_id
                try:
                    zluda_installer.load(zluda_path)
                    torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.0 torchvision --index-url https://download.pytorch.org/whl/cu118')