nikigoli commited on
Commit
768648c
1 Parent(s): 41359b1

Update models/GroundingDINO/ops/setup.py

Browse files
Files changed (1) hide show
  1. models/GroundingDINO/ops/setup.py +3 -0
models/GroundingDINO/ops/setup.py CHANGED
@@ -11,6 +11,8 @@ import glob
11
 
12
  import torch
13
 
 
 
14
  from torch.utils.cpp_extension import CUDA_HOME
15
  from torch.utils.cpp_extension import CppExtension
16
  from torch.utils.cpp_extension import CUDAExtension
@@ -20,6 +22,7 @@ from setuptools import setup
20
 
21
  requirements = ["torch", "torchvision"]
22
 
 
23
  def get_extensions():
24
  this_dir = os.path.dirname(os.path.abspath(__file__))
25
  extensions_dir = os.path.join(this_dir, "src")
 
11
 
12
  import torch
13
 
14
+ import spaces
15
+
16
  from torch.utils.cpp_extension import CUDA_HOME
17
  from torch.utils.cpp_extension import CppExtension
18
  from torch.utils.cpp_extension import CUDAExtension
 
22
 
23
  requirements = ["torch", "torchvision"]
24
 
25
+ @spaces.GPU
26
  def get_extensions():
27
  this_dir = os.path.dirname(os.path.abspath(__file__))
28
  extensions_dir = os.path.join(this_dir, "src")