Update models/GroundingDINO/ops/setup.py
Browse files
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")
|