eliphatfs commited on
Commit
ba37554
1 Parent(s): 43ad84a

Fix for categories of different lengths.

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -154,7 +154,7 @@ def demo_classification():
154
  col2 = misc_utils.render_pc(pc)
155
  prog.progress(0.5, "Computing Category Embeddings")
156
  device = clip_model.device
157
- tn = clip_prep(text=cats, return_tensors='pt', truncation=True, max_length=76).to(device)
158
  feats = clip_model.get_text_features(**tn).float().cpu()
159
  prog.progress(0.5, "Running Classification")
160
  pred = classification.pred_custom_sims(model_g14, pc, cats, feats)
 
154
  col2 = misc_utils.render_pc(pc)
155
  prog.progress(0.5, "Computing Category Embeddings")
156
  device = clip_model.device
157
+ tn = clip_prep(text=cats, return_tensors='pt', truncation=True, max_length=76, padding=True).to(device)
158
  feats = clip_model.get_text_features(**tn).float().cpu()
159
  prog.progress(0.5, "Running Classification")
160
  pred = classification.pred_custom_sims(model_g14, pc, cats, feats)