fix n_gpu count when no_cuda flag is activated (#3077)

* fix n_gpu count when no_cuda flag is activated * someone was left behind
2020-03-02 10:20:21 -05:00
parent 298bed16a8
commit 6b1ff25084
13 changed files with 13 additions and 13 deletions
--- a/examples/run_generation.py
+++ b/examples/run_generation.py
@@ -189,7 +189,7 @@ def main():
    args = parser.parse_args()

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
-    args.n_gpu = torch.cuda.device_count()
+    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()

    set_seed(args)