taking care if cuda fails

montrealrobotics · Sep 18, 2023 · b7e9ca3 · b7e9ca3
1 parent 2f6a2b3
commit b7e9ca3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cleanrl/ppo_continuous_action_wandb.py b/cleanrl/ppo_continuous_action_wandb.py
@@ -446,7 +446,7 @@ def train(cfg):
     torch.manual_seed(cfg.model_seed)
     torch.backends.cudnn.deterministic = cfg.torch_deterministic
 
-    device = torch.device("cuda" if torch.cuda.is_available() and cfg.cuda else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 and cfg.cuda else "cpu")
 
     # env setup
     envs = gym.vector.SyncVectorEnv(