fix(yzj):fix device bug

opendilab · Nov 22, 2023 · 8d71f96 · 8d71f96
1 parent 6ea3f9b
commit 8d71f96
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/lzero/policy/efficientzero.py b/lzero/policy/efficientzero.py
@@ -401,6 +401,7 @@ def _forward_learn(self, data: torch.Tensor) -> Dict[str, Union[float, int]]:
                 # obtain the oracle latent states from representation function.
                 beg_index, end_index = self._get_target_obs_index_in_step_k(step_k)
                 obs_target_batch_tmp = default_collate(obs_target_batch[:, beg_index:end_index].squeeze())
+                obs_target_batch_tmp = to_device(obs_target_batch_tmp, self._device)
                 network_output = self._learn_model.initial_inference(obs_target_batch_tmp)
 
                 latent_state = to_tensor(latent_state)

diff --git a/lzero/policy/muzero.py b/lzero/policy/muzero.py
@@ -369,6 +369,7 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
                     # obtain the oracle latent states from representation function.
                     beg_index, end_index = self._get_target_obs_index_in_step_k(step_k)
                     obs_target_batch_tmp = default_collate(obs_target_batch[:, beg_index:end_index].squeeze())
+                    obs_target_batch_tmp = to_device(obs_target_batch_tmp, self._device)
                     network_output = self._learn_model.initial_inference(obs_target_batch_tmp)
 
                     latent_state = to_tensor(latent_state)