@torch.no_grad()
def predict_pi(self, x: torch.Tensor) -> np.ndarray:
    self.eval()
    (logits, _) = self._get_dist_params(x)
    self.train()
    return F.softmax(logits, dim=(- 1)).detach().cpu().numpy()
@torch.no_grad()
def predict_pi(self, x: torch.Tensor) -> np.ndarray:
    self.eval()
    (logits, _) = self._get_dist_params(x)
    self.train()
    return F.softmax(logits, dim=(- 1)).detach().cpu().numpy()