code: def test_model_train_output(self, made, demo_data):
      made.train()
      out = made(demo_data["data"])
  
      assert isinstance(out, ModelOutput)
  
      assert set(["mu", "log_var"]) == set(out.keys())
  
      assert out.mu.shape[0] == demo_data["data"].shape[0]
      assert out.log_var.shape[0] == demo_data["data"].shape[0]
      assert out.mu.shape[1:] == np.prod(made.model_config.output_dim)
      assert out.log_var.shape[1:] == np.prod(
          made.model_config.output_dim
      )
  
file_path: benchmark_VAE/tests/test_MADE.py
imports: import os
  import numpy
  import pytest
  import torch
  from pythae.models import AutoModel
  from pythae.models.base.base_utils import ModelOutput
  from pythae.models.normalizing_flows import MADE, MADEConfig
label_desc: Use image node_id 3 for calling the Test_Model_forward obj's underlying member method code with example usage: obj.test_model_train_output(made, demo_data) without return types