yo

mosaicml · Aug 30, 2024 · 4a96d02 · 4a96d02
1 parent 94bc894
commit 4a96d02
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -128,12 +128,12 @@
 
 extra_deps['databricks-serverless'] = {
     dep for key, deps in extra_deps.items() for dep in deps
-    if 'gpu' not in key and 'megablocks' not in key and
+    if 'gpu' not in key and 'megablocks' not in key and 'te' not in key and
     'databricks-connect' not in dep
 }
 extra_deps['all-cpu'] = {
     dep for key, deps in extra_deps.items() for dep in deps
-    if 'gpu' not in key and 'megablocks' not in key
+    if 'gpu' not in key and 'megablocks' not in key and 'te' not in key
 }
 extra_deps['all'] = {
     dep for key, deps in extra_deps.items() for dep in deps