Add Accelerate for GHCI

baal-org · Apr 4, 2024 · b608a08 · b608a08
1 parent d64e01e
commit b608a08
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 3 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -32,6 +32,7 @@ lightning-flash = { version = ">=0.7.5", optional=true }
 
 # NLP
 transformers = {version = ">=4.10.2", optional=true}
+accelerate = {version = "^0.28.0", optional=true}
 datasets = {version = ">=1.11.0", optional=true}
 
 [tool.poetry.dev-dependencies]
@@ -57,7 +58,7 @@ mkdocs-exclude-search = "^0.6.4"
 
 [tool.poetry.extras]
 vision = ["torchvision", "lightning-flash"]
-nlp = ["transformers", "datasets"]
+nlp = ["transformers", "datasets", "accelerate"]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]