code: def test_multi_input_model():
      class MultiInModel(nn.Module):
          def __init__(self):
              super().__init__()
              self.model = SimpleModel()
  
          def forward(self, x):
              # We get two inputs
              x1, x2 = x
              # We merge those inputs
              return self.model(x1) + self.model(x2)
  
      model = MultiInModel()
      wrapper = ModelWrapper(model, None)
      dataset = DummyDataset(n_in=2)
      assert len(dataset[0]) == 2
      b = next(iter(DataLoader(dataset, 15, False)))[0]
      l = wrapper.predict_on_batch(b, iterations=10, cuda=False)
      assert l.shape[0] == 15 and l.shape[-1] == 10
  
imports: import math
  import unittest
  from unittest.mock import Mock
  import numpy
  import pytest
  import torch
  import torchmetrics
  from torch import nn
  from torch.utils.data import Dataset, DataLoader
  from baal.modelwrapper import ModelWrapper, mc_inference
  from baal.utils.metrics import ClassificationReport
label_desc: Use image node_id 1 for calling a global function with example usage: test_multi_input_model() and returns: unknown