From a872a6283f172438594b20f56267da8ad347da29 Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Mon, 28 Mar 2022 14:25:33 +0800
Subject: [PATCH] Update init

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 __init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/__init__.py b/__init__.py
index 85a7c81..2e83d9a 100644
--- 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)