Slep commited on
Commit
aeb008e
1 Parent(s): ba56501

Upload CondViTForEmbedding

Browse files
Files changed (2) hide show
  1. config.json +1 -1
  2. hf_model.py +2 -2
config.json CHANGED
@@ -1,5 +1,5 @@
1
  {
2
- "_name_or_path": "CondViT-B16-cat",
3
  "architectures": [
4
  "CondViTForEmbedding"
5
  ],
 
1
  {
2
+ "_name_or_path": "__debug_save",
3
  "architectures": [
4
  "CondViTForEmbedding"
5
  ],
hf_model.py CHANGED
@@ -43,5 +43,5 @@ class CondViTForEmbedding(PreTrainedModel):
43
  n_categories=config.n_categories,
44
  )
45
 
46
- def forward(self, img_tensors, category_indices=None):
47
- return self.model(img_tensors, category_indices)
 
43
  n_categories=config.n_categories,
44
  )
45
 
46
+ def forward(self, pixel_values, category_indices=None):
47
+ return self.model(imgs=pixel_values, c=category_indices)