code: def test_return_sampling(
      self, model, dummy_data, sampler, num_sample_and_batch_size
  ):
      num_samples, batch_size = (
          num_sample_and_batch_size[0],
          num_sample_and_batch_size[1],
      )
  
      sampler.fit(train_data=dummy_data)
  
      gen_samples = sampler.sample(
          num_samples=num_samples,
          batch_size=batch_size,
          return_gen=True,
      )
  
      assert gen_samples.shape[0] == num_samples
  
file_path: benchmark_VAE/tests/test_vamp_sampler.py
imports: import os
  import numpy
  import pytest
  import torch
  from pythae.models import VAMP, VAMPConfig
  from pythae.pipelines import GenerationPipeline
  from pythae.samplers import NormalSampler, NormalSamplerConfig, VAMPSampler, VAMPSamplerConfig
label_desc: Use image node_id 1 for calling the Test_VAMPSampler_Sampling obj's underlying member method code with example usage: obj.test_return_sampling(model, dummy_data, sampler, num_sample_and_batch_size) without return types