gpu setup

inference/inference.py

@@ -77,7 +77,7 @@ def get_heatmap_rgb(model, data, N_CHUNKS=5): # evaluate loader can only have ba
         heatmap_rgb = torch.tensor(plt.cm.jet(scores.numpy())[:,:3]).squeeze()
         return heatmap_rgb
     
-@spaces.GPU(duration=90)
+@spaces.GPU
 def segment_obj(xyz, rgb, normal, queries):
     model = load_model()
     data_dict = preprocess_pcd(torch.tensor(xyz).float().to(DEVICE), torch.tensor(rgb).float().to(DEVICE), torch.tensor(normal).float().to(DEVICE))
@@ -85,7 +85,7 @@ def segment_obj(xyz, rgb, normal, queries):
     seg_rgb = get_segmentation_rgb(model, data_dict)
     return seg_rgb
     
-@spaces.GPU(duration=90)
+@spaces.GPU
 def get_heatmap(xyz, rgb, normal, query):
     model = load_model()
     data_dict = preprocess_pcd(torch.tensor(xyz).float().to(DEVICE), torch.tensor(rgb).float().to(DEVICE), torch.tensor(normal).float().to(DEVICE))

requirements.txt

@@ -24,7 +24,7 @@ scipy
 plyfile
 termcolor
 timm
-spconv
+spconv-cu118
 transformers
 open3d
 sentencepiece