brozek/rltorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Attempting to see if this fixes the mismatched devices error

Browse source
This commit is contained in:
Brandon Rozek 2019-02-03 00:49:47 -05:00
parent 9cd3625fd3
commit 03455accc8
1 changed files with 7 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
rltorch/agents/DQNAgent.py
View file

@ -17,6 +17,13 @@ class DQNAgent:
minibatch = self.memory.sample(self.config['batch_size'])
state_batch, action_batch, reward_batch, next_state_batch, not_done_batch = M.zip_batch(minibatch)
# Send to their appropriate devices
state_batch = state_batch.to(net.device)
action_batch = action_batch.to(net.device)
reward_batch = reward_batch.to(net.device)
next_state_batch = next_state_batch.to(net.device)
not_done_batch = not_done_batch.to(net.device)
obtained_values = self.net(state_batch).gather(1, action_batch.view(self.config['batch_size'], 1))