Update app.py

app.py

@@ -107,7 +107,7 @@ def classification_custom(load_data, cats):
     ref_dev = next(model_g14.parameters()).device
     enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev))
     if model_name == "pb-sn-M":
-        enc = pc_adapter(enc)    
+        enc = pc_adapter(enc)  
     sim = torch.matmul(torch.nn.functional.normalize(feats, dim=-1), torch.nn.functional.normalize(enc.cpu(), dim=-1).squeeze())
     argsort = torch.argsort(sim, descending=True)
     pred = OrderedDict((cats[i], sim[i]) for i in argsort if i < len(cats))
@@ -123,9 +123,10 @@ def retrieval_pc(load_data, k, sim_th, filter_fn):
     prog.progress(0.5, "Computing Embeddings")
     col2 = utils.render_pc(pc)
     ref_dev = next(model_g14.parameters()).device
-    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
-
-    sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc, dim=-1).squeeze())
+    enc = model_g14(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev))
+    if model_name == "pb-sn-M":
+        enc = pc_adapter(enc)
+    sim = torch.matmul(torch.nn.functional.normalize(lvis.feats, dim=-1), torch.nn.functional.normalize(enc.cpu(), dim=-1).squeeze())
     argsort = torch.argsort(sim, descending=True)
     pred = OrderedDict((lvis.categories[i], sim[i]) for i in argsort if i < len(lvis.categories))
     with col2:
@@ -213,7 +214,7 @@ try:
     if model_name == "pb-sn-M":
         model_g14, pc_adapter = load_tripletmix('tripletmix-pointbert-shapenet')
     elif model_name == "pb-sn":
-        model_g14 = load_openshape('openshape-pointbert-vitg14-rgb')
+        model_g14 = load_openshape('openshape-pointbert-shapenet')
     task = st.sidebar.selectbox(
         'Task Selection', 
         ("3D Classification", "Cross-modal retrieval", "Cross-modal generation")