Improve: Use the device_map option in transformers to load BERT models directly to the GPU

2024-08-10 16:56:19 +09:00
parent bdf20e8911
commit 80dd6dbc22
7 changed files with 42 additions and 20 deletions
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,6 +22,7 @@ classifiers = [
    "Programming Language :: Python :: Implementation :: CPython",
 ]
 dependencies = [
+    "accelerate",
    "cmudict",
    "cn2an",
    "g2p_en",