Zhen Ye commited on
Commit
fc53373
·
1 Parent(s): 9606129

fix:multi-gpu support

Browse files
Files changed (1) hide show
  1. 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