wassname/DeepRL
1
0
Fork 0
mirror of https://github.com/wassname/DeepRL.git synced 2026-09-09 11:13:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix a bug for GPU

Browse Source
This commit is contained in:
Shangtong Zhang
2018-04-25 17:01:05 -06:00
parent 6f517e3eef
commit 4f2903eb8f
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
network/network_utils.py
+1 -1
View File
@@ -12,7 +12,7 @@ import numpy as np
class BaseNet:
def set_gpu(self, gpu):
if gpu >= 0 and torch.cuda.is_available():
self.device = torch.device('gpu:%d' % (gpu))
self.device = torch.device('cuda:%d' % (gpu))
else:
self.device = torch.device('cpu')
self.to(self.device)