Update openshape/__init__.py

openshape/__init__.py

@@ -32,6 +32,13 @@ def B32(s):
     model.load_state_dict(module(s, 'pc_encoder'))
     return model
 
+def S14(s):
+    model = Projected(
+        PointPatchTransformer(256, 6, 4, 1024, 96, 64, 0.4, 256, 6),
+        nn.Linear(256, cfg.model.out_channel)
+    )
+    model.load_state_dict(module(['state_dict'], 'module'))
+    return model
 
 def G14_M(s):
     model = Projected(
@@ -50,6 +57,7 @@ model_list = {
     "openshape-pointbert-vitb32-rgb": B32,
     "openshape-pointbert-vitl14-rgb": L14,
     "openshape-pointbert-vitg14-rgb": G14,
+    "openshape-pointbert-shapenet": S14,
     "tripletmix-pointbert-shapenet": G14_M,
     "pc_adapter": PCA,
 }