Update Categorical_DQN.py

fix map_location
XinJingHao · Jun 8, 2024 · a464e0a · a464e0a
1 parent 987e83b
commit a464e0a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/2.4_Categorical-DQN_C51/Categorical_DQN.py b/2.4_Categorical-DQN_C51/Categorical_DQN.py
@@ -114,8 +114,8 @@ def save(self,algo,EnvName,steps):
 		torch.save(self.q_net.state_dict(), "./model/{}_{}_{}k.pth".format(algo,EnvName,steps))
 
 	def load(self,algo,EnvName,steps):
-		self.q_net.load_state_dict(torch.load("./model/{}_{}_{}k.pth".format(algo,EnvName,steps)))
-		self.q_target.load_state_dict(torch.load("./model/{}_{}_{}k.pth".format(algo,EnvName,steps)))
+		self.q_net.load_state_dict(torch.load("./model/{}_{}_{}k.pth".format(algo,EnvName,steps), map_location=self.dvc))
+		self.q_target.load_state_dict(torch.load("./model/{}_{}_{}k.pth".format(algo,EnvName,steps), map_location=self.dvc))
 
 
 class ReplayBuffer(object):