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)