mirror of
https://github.com/open-mmlab/mmengine.git
synced 2025-06-03 21:54:44 +08:00
[Enhancement] Handle the case for Multi-Instance GPUs when using cuda_visible_devices (#1164)
This commit is contained in:
parent
691500dce6
commit
0ff89f7c36
@ -376,7 +376,12 @@ def _get_device_id():
|
|||||||
cuda_visible_devices = list(range(num_device))
|
cuda_visible_devices = list(range(num_device))
|
||||||
else:
|
else:
|
||||||
cuda_visible_devices = cuda_visible_devices.split(',')
|
cuda_visible_devices = cuda_visible_devices.split(',')
|
||||||
|
if cuda_visible_devices[local_rank].isdigit():
|
||||||
return int(cuda_visible_devices[local_rank])
|
return int(cuda_visible_devices[local_rank])
|
||||||
|
else:
|
||||||
|
# handle case for Multi-Instance GPUs
|
||||||
|
# see #1148 for details
|
||||||
|
return cuda_visible_devices[local_rank]
|
||||||
|
|
||||||
|
|
||||||
def _get_host_info() -> str:
|
def _get_host_info() -> str:
|
||||||
|
Loading…
x
Reference in New Issue
Block a user