code: def on_validation_epoch_end(self) -> None:
      pi = [
          nn.functional.gumbel_softmax(
              output[2]["pi"],
              tau=self.head.hparams.softmax_temperature,
              dim=-1,
          )
          for output in self._val_output
      ]
      pi = torch.cat(pi).detach().cpu()
      for i in range(self.head.hparams.num_gaussian):
          self.log(
              f"mean_pi_{i}",
              pi[:, i].mean(),
              on_epoch=True,
              on_step=False,
              logger=True,
              prog_bar=False,
          )
  
      mu = [output[2]["mu"] for output in self._val_output]
      mu = torch.cat(mu).detach().cpu()
      for i in range(self.head.hparams.num_gaussian):
          self.log(
              f"mean_mu_{i}",
              mu[:, i].mean(),
              on_epoch=True,
              on_step=False,
              logger=True,
              prog_bar=False,
          )
  
      sigma = [output[2]["sigma"] for output in self._val_output]
      sigma = torch.cat(sigma).detach().cpu()
      for i in range(self.head.hparams.num_gaussian):
          self.log(
              f"mean_sigma_{i}",
              sigma[:, i].mean(),
              on_epoch=True,
              on_step=False,
              logger=True,
              prog_bar=False,
          )
      if self.do_log_logits:
          logits = [output[0] for output in self._val_output]
          logits = torch.cat(logits).detach().cpu()
          fig = self.create_plotly_histogram(
              logits.unsqueeze(1), "logits"
          )
          wandb.log(
              {
                  "valid_logits": fig,
                  "global_step": self.global_step,
              },
              commit=False,
          )
          if self.head.hparams.log_debug_plot:
              fig = self.create_plotly_histogram(
                  pi,
                  "pi",
                  bin_dict={"start": 0.0, "end": 1.0, "size": 0.1},
              )
              wandb.log(
                  {
                      "valid_pi": fig,
                      "global_step": self.global_step,
                  },
                  commit=False,
              )
  
              fig = self.create_plotly_histogram(mu, "mu")
              wandb.log(
                  {
                      "valid_mu": fig,
                      "global_step": self.global_step,
                  },
                  commit=False,
              )
  
              fig = self.create_plotly_histogram(sigma, "sigma")
              wandb.log(
                  {
                      "valid_sigma": fig,
                      "global_step": self.global_step,
                  },
                  commit=False,
              )
      self._val_output = []
  
imports: from typing import Dict, Optional, Union
  import torch
  import torch.nn
  from omegaconf import DictConfig, OmegaConf
  from torch import Tensor
  from pytorch_tabular import models
  from pytorch_tabular.config.config import ModelConfig
  from pytorch_tabular.models.common.heads import blocks
  from pytorch_tabular.models.tab_transformer.tab_transformer import TabTransformerBackbone
  from pytorch_tabular.tabular_model import getattr_nested
  from pytorch_tabular.utils import get_logger
  from ..base_model import BaseModel, safe_merge_config
label_desc: Use image node_id 17 for calling the MDNModel obj's underlying member method code with example usage: obj.on_validation_epoch_end() without return types