1.5T with more CUDA memory for plan A

Tencent · Apr 30, 2021 · eada68d · eada68d
1 parent 5f67e4d
commit eada68d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_bert.py b/tests/test_bert.py
@@ -203,7 +203,7 @@ def test_bert_model(is_ckp: bool = False,
 
     if is_ps:
         manager = HybridPSManager()
-        manager.init([1024 * 1024 * 1024 * 2] * 1,
+        manager.init([1024 * 1024 * 1024 * 8] * 1,
                      [1024 * 1024 * 1024 * 4 * 4])
         # chunk 512 MB, good for CPU-GPU bandwidth
         client = HybridPSClient(gpu_index=0,