code: def test_prediction(use_cuda, n_ensemble):
      model = AModel()
      ensemble = EnsembleModelWrapper(model, nn.CrossEntropyLoss())
      optimizer = optim.SGD(model.parameters(), lr=0.001)
      dataset = DummyDataset()
      if use_cuda:
          model.cuda()
  
      for i in range(n_ensemble):
          model.apply(weight_init)
          ensemble.train_on_dataset(dataset, optimizer, 1, 2, use_cuda)
          ensemble.add_checkpoint()
      assert len(ensemble._weights) == n_ensemble
  
      out = ensemble.predict_on_batch(
          dataset[0][0].unsqueeze(0), cuda=use_cuda
      )
      assert not all(
          torch.eq(out[..., 0], out[..., i]).all()
          for i in range(1, n_ensemble)
      )
  
      out = ensemble_prediction(
          model=model,
          data=dataset[1][0].unsqueeze(0),
          weights=ensemble._weights,
          cuda=use_cuda,
      )
      assert not all(
          torch.eq(out[..., 0], out[..., i]).all()
          for i in range(1, n_ensemble)
      )
  
      ensemble.clear_checkpoints()
      assert len(ensemble._weights) == 0
  
      with pytest.raises(ValueError):
          ensemble.predict_on_batch(
              dataset[0][0].unsqueeze(0), cuda=use_cuda
          )
  
file_path: baal/tests/ensemble_test.py
imports: import pytest
  import torch
  from torch import nn
  from torch import optim
  from torch.utils.data import Dataset
  from baal.ensemble import EnsembleModelWrapper, ensemble_prediction
label_desc: Use image node_id 2 for calling a global function with example usage: test_prediction(use_cuda, n_ensemble) without return types