zerogpu

inference/inference.py

@@ -4,9 +4,9 @@ import matplotlib.pyplot as plt
 from inference.utils import get_seg_color, load_model, preprocess_pcd, encode_text
 import spaces
 
-DEVICE = "cpu"
-if torch.cuda.is_available():
-    DEVICE = "cuda:0"
+DEVICE = "cuda:0"
+#if torch.cuda.is_available():
+    #DEVICE = "cuda:0"
 
 def pred_3d_upsample(
         pred, # n_subsampled_pts, feat_dim

inference/utils.py

@@ -5,9 +5,9 @@ import numpy as np
 import random
 from transformers import AutoTokenizer, AutoModel
 
-DEVICE = "cpu"
-if torch.cuda.is_available():
-    DEVICE = "cuda:0"
+DEVICE = "cuda:0"
+#if torch.cuda.is_available():
+    #DEVICE = "cuda:0"
 
 def get_seg_color(labels):
     part_num = labels.max()