Update init

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
3 years ago · a872a6283f
1 changed files with 2 additions and 2 deletions
--- a/init.py
+++ b/init.py
@ -15,5 +15,5 @@
 from .timm_image import TimmImage


-def timm():
-    return TimmImage()
+def timm(model_name: str, num_classes: int=1000):
+    return TimmImage(model_name=model_name, num_classes=num_classes)