updated args

2019-06-25 18:23:29 -04:00 · 2019-06-25 18:23:29 -04:00 · 338f889e7c
parent e12c8ad21a
commit 338f889e7c
1 changed files with 1 additions and 3 deletions
--- a/docs/source/examples/fully_featured_trainer.py
+++ b/docs/source/examples/fully_featured_trainer.py
@ -43,9 +43,7 @@ def main(hparams, cluster, results_dict):
    :param hparams:
    :return:
    """
-    on_gpu = torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
+    on_gpu = hparams.gpus is not None and torch.cuda.is_available()

    device = 'cuda' if on_gpu else 'cpu'
    hparams.__setattr__('device', device)