Update difpoint/src/models/predictor.py

difpoint/src/models/predictor.py

@@ -10,6 +10,7 @@ import torch
 from torch.cuda import nvtx
 from collections import OrderedDict
 import platform
+import spaces
 
 try:
     import tensorrt as trt
@@ -176,7 +177,7 @@ class OnnxRuntimePredictor:
     """
     OnnxRuntime Prediction
     """
-
+    @spaces.GPU
     def __init__(self, **kwargs):
         model_path = kwargs.get("model_path", "")  # 用模型路径区分是否是一样的实例
         assert os.path.exists(model_path), "model path must exist!"