Fixed error for CUDA

This commit is contained in:
Vasilis Valatsos 2024-03-14 20:14:16 +02:00
parent 9e963c1ab3
commit 9057f070cf

View file

@ -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.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) 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.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) self.to(self.device)