make multi tensor model onnx compatible (again)

Author: FynnBe

example_descriptions/models/unet2d_multi_tensor/bioimageio.yaml

@@ -143,7 +143,7 @@ weights:
         in_channels: 2
         initial_features: 16
         out_channels: 2
-      sha256: 89204f8f3513b3c227127a8137bedaec9eafe49925f7734c73c6650ec135b34e
+      sha256: 74b6d27cd17b40560e70fde4d57b5675614979ffab4f1e0328a6c3a3f64a1ff2
       source: multi_tensor_unet.py
     pytorch_version: 1.6
     sha256: c498522b3f2b02429b41fe9dbcb722ce0d7ad4cae7fcf8059cee27857ae49b00

example_descriptions/models/unet2d_multi_tensor/multi_tensor_unet.py

@@ -1,4 +1,6 @@
 # type: ignore
+from typing import List, Optional
+
 import torch
 import torch.nn as nn
 
@@ -75,11 +77,16 @@ def _apply_default(self, x):
 
         return x
 
-    def forward(self, *x):
-        assert isinstance(x, (list, tuple)), type(x)
-        # fix issue in onnx export
-        if isinstance(x[0], list) and len(x) == 1:
-            x = x[0]
+    def forward(
+        self,
+        x0: torch.Tensor,
+        x1: Optional[torch.Tensor] = None,
+        x2: Optional[torch.Tensor] = None,
+        x3: Optional[torch.Tensor] = None,
+        x4: Optional[torch.Tensor] = None,
+        /,
+    ) -> List[torch.Tensor]:
+        x = [x for x in [x0, x1, x2, x3, x4] if x is not None]
         assert len(x) == self.in_channels, f"{len(x)}, {self.in_channels}"
         x = torch.cat(x, dim=1)
         out = self._apply_default(x)