Spaces:
Running
on
A10G
Running
on
A10G
Zhen Ye
commited on
Commit
·
fc53373
1
Parent(s):
9606129
fix:multi-gpu support
Browse files- inference.py +8 -0
inference.py
CHANGED
|
@@ -5,6 +5,14 @@ from threading import RLock, Thread
|
|
| 5 |
from queue import Queue, PriorityQueue
|
| 6 |
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
| 7 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
import cv2
|
| 9 |
import numpy as np
|
| 10 |
import torch
|
|
|
|
| 5 |
from queue import Queue, PriorityQueue
|
| 6 |
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
| 7 |
|
| 8 |
+
# Must clear CUDA_VISIBLE_DEVICES before importing torch to ensure all GPUs are seen
|
| 9 |
+
if "CUDA_VISIBLE_DEVICES" in os.environ:
|
| 10 |
+
# We only clear it if we are sure we want all GPUs.
|
| 11 |
+
# In HF Spaces 4xGPU, this var might be set to "0" by default?
|
| 12 |
+
# Let's log it first implies we need imports.
|
| 13 |
+
# Just Unset it.
|
| 14 |
+
del os.environ["CUDA_VISIBLE_DEVICES"]
|
| 15 |
+
|
| 16 |
import cv2
|
| 17 |
import numpy as np
|
| 18 |
import torch
|