Correct behavior for argument gpus in Trainer (#561)

Lightning-AI · Nov 30, 2019 · df7b6d9 · df7b6d9
1 parent db0587f
commit df7b6d9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
@@ -132,7 +132,7 @@ def __init__(self,
         self.gradient_clip_val = gradient_clip_val
         self.check_val_every_n_epoch = check_val_every_n_epoch
         self.track_grad_norm = track_grad_norm
-        self.on_gpu = gpus is not None and torch.cuda.is_available()
+        self.on_gpu = True if (gpus and torch.cuda.is_available()) else False
         self.process_position = process_position
         self.weights_summary = weights_summary
         self.max_nb_epochs = max_nb_epochs