From 9057f070cf566917f78a573d1a32f03f17de6561 Mon Sep 17 00:00:00 2001 From: Vasilis Valatsos Date: Thu, 14 Mar 2024 20:14:16 +0200 Subject: [PATCH] Fixed error for CUDA --- ml/ppo/brain.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ml/ppo/brain.py b/ml/ppo/brain.py index d522e4a..75153da 100644 --- a/ml/ppo/brain.py +++ b/ml/ppo/brain.py @@ -70,7 +70,7 @@ class ActorNetwork(nn.Module): self.optimizer = optim.Adam(self.parameters(), lr=alpha, betas=(0.9, 0.9), eps=1e-5) - self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu') + self.device = T.device('cuda' if T.cuda.is_available() else 'cpu') self.to(self.device) @@ -117,7 +117,7 @@ class CriticNetwork(nn.Module): self.optimizer = optim.Adam(self.parameters(), lr=alpha, betas=(0.9, 0.9), eps=1e-5) - self.device = T.device('cuda:1' if T.cuda.is_available() else 'cpu') + self.device = T.device('cuda' if T.cuda.is_available() else 'cpu') self.to(self.device)