[s2s] test_distributed_eval (#8315)

Co-authored-by: Sam Shleifer <sshleifer@gmail.com>
2020-11-05 13:01:15 -08:00
parent 04e442d575
commit d787935a14
4 changed files with 56 additions and 8 deletions
--- a/docs/source/testing.rst
+++ b/docs/source/testing.rst
@@ -450,7 +450,8 @@ Inside tests:

 .. code-block:: bash

-   torch.cuda.device_count()
+   from transformers.testing_utils import get_gpu_count
+   n_gpu = get_gpu_count() # works with torch and tf