Update openshape/__init__.py

openshape/__init__.py

@@ -58,6 +58,7 @@ model_list = {
     "openshape-pointbert-vitl14-rgb": L14,
     "openshape-pointbert-vitg14-rgb": G14,
     "openshape-pointbert-shapenet": S14,
+    "tripletmix-pointbert-all-modelnet40": G14_M,
     "tripletmix-pointbert-shapenet": G14_M,
     "pc_adapter": PCA,
 }
@@ -72,13 +73,8 @@ def load_pc_encoder(name):
 
 
 def load_pc_encoder_mix(name):
-    pc_adapter = None
     s = torch.load(hf_hub_download("TripletMix/" + name, "model.pt"), map_location='cpu')
     model = model_list[name](s).eval()
-    if name ==  "tripletmix-pointbert-shapenet":
-        pc_adapter = model_list["pc_adapter"](s).eval()
-        if torch.cuda.is_available():
-            pc_adapter.cuda()
     if torch.cuda.is_available():
         model.cuda()
-    return model, pc_adapter
+    return model