fix(yzj):fix device bug

Author: xindong.he

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)

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)