code: def test_CEVAE():
      y, X, treatment, tau, b, e = simulate_hidden_confounder(
          n=10000, p=5, sigma=1.0, adj=0.0
      )
  
      outcome_dist = "normal"
      latent_dim = 20
      hidden_dim = 200
      num_epochs = 50
      batch_size = 100
      learning_rate = 1e-3
      learning_rate_decay = 0.1
  
      cevae = CEVAE(
          outcome_dist=outcome_dist,
          latent_dim=latent_dim,
          hidden_dim=hidden_dim,
          num_epochs=num_epochs,
          batch_size=batch_size,
          learning_rate=learning_rate,
          learning_rate_decay=learning_rate_decay,
      )
  
      cevae.fit(
          X=torch.tensor(X, dtype=torch.float),
          treatment=torch.tensor(treatment, dtype=torch.float),
          y=torch.tensor(y, dtype=torch.float),
      )
  
      # check the accuracy of the ite accuracy
      ite = cevae.predict(X).flatten()
  
      auuc_metrics = pd.DataFrame(
          {
              "ite": ite,
              "W": treatment,
              "y": y,
              "treatment_effect_col": tau,
          }
      )
  
      cumgain = get_cumgain(
          auuc_metrics,
          outcome_col="y",
          treatment_col="W",
          treatment_effect_col="tau",
      )
  
      # Check if the cumulative gain when using the model's prediction is
      # higher than it would be under random targeting
      assert cumgain["ite"].sum() > cumgain["Random"].sum()
  
imports: import pandas
  import torch
  from causalml.inference.nn import CEVAE
  from causalml.dataset import simulate_hidden_confounder
  from causalml.metrics import get_cumgain
label_desc: Use image node_id 1 for calling a global function with example usage: test_CEVAE() without return types