sravanneeli commited on
Commit
44ba809
·
1 Parent(s): 0d3cb67

Add new models for image classification

Browse files
Files changed (1) hide show
  1. vision_models.py +4 -1
vision_models.py CHANGED
@@ -5,7 +5,10 @@ import numpy as np
5
 
6
  classification_models = {
7
  "ResNet18": "resnet_18_imagenet",
8
- "ResNet50": "resnet_50_imagenet"
 
 
 
9
  }
10
 
11
 
 
5
 
6
  classification_models = {
7
  "ResNet18": "resnet_18_imagenet",
8
+ "ResNet50": "resnet_50_imagenet",
9
+ "ViT-B16-224": "vit_base_patch16_224_imagenet",
10
+ "ViT-B16-384": "vit_base_patch16_384_imagenet",
11
+ "ViT-L16-224": "vit_large_patch16_224_imagenet"
12
  }
13
 
14