Update README.md

README.md

@@ -32,7 +32,7 @@ transform = ImageNetTransform(train=False, img_size=512)
 x = transform(read_rgb("cat.png"))
 x = mx.expand_dims(x, 0)
 
-model = create_model("vit_large_patch16_512.swag_e2e-mlxim")
+model = create_model("vit_large_patch16_512.swag_e2e")
 model.eval()
 
 logits = model(x)
@@ -49,16 +49,16 @@ x = transform(read_rgb("cat.png"))
 x = mx.expand_dims(x, 0)
 
 # first option
-model = create_model("vit_large_patch16_512.swag_e2e-mlxim", num_classes=0)
+model = create_model("vit_large_patch16_512.swag_e2e", num_classes=0)
 model.eval()
 
 embeds = model(x)
 
 # second option
-model = create_model("vit_large_patch16_512.swag_e2e-mlxim")
+model = create_model("vit_large_patch16_512.swag_e2e")
 model.eval()
 
-embeds = model.features(x)
+embeds = model.get_features(x)
 ```