Add model_utils

2023-02-13 09:36:12 +00:00
parent bb59a3e730
commit fffa2e1f4b
2 changed files with 17 additions and 2 deletions
--- a/cacheflow/worker/models/init.py
+++ b/cacheflow/worker/models/init.py
@@ -1,5 +1,7 @@
-from cacheflow.worker.models.opt import OPTForCausalLM
+from cacheflow.worker.models.model_utils import get_model
+

 __all__ = [
-    'OPTForCausalLM',
+    'get_model',
+    
 ]