Fixed error for CUDA
This commit is contained in:
parent
9e963c1ab3
commit
9057f070cf
1 changed files with 2 additions and 2 deletions
|
@ -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)
|
||||
|
||||
|
|
Loading…
Reference in a new issue