From 338f889e7c1d114798d37fce2533a1c212f37bce Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 25 Jun 2019 18:23:29 -0400 Subject: [PATCH] updated args --- docs/source/examples/fully_featured_trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/source/examples/fully_featured_trainer.py b/docs/source/examples/fully_featured_trainer.py index dfb90bb11b..4037ade812 100644 --- 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)