diff --git a/training/training_scripts/NN_Filtering/HOP4/model/model.py b/training/training_scripts/NN_Filtering/HOP4/model/model.py index 4e484407e23f6c0dac345f0355e5c351f2c5019b..76c7b698f55684eb7c4bb689eca46b98b22927a1 100644 --- a/training/training_scripts/NN_Filtering/HOP4/model/model.py +++ b/training/training_scripts/NN_Filtering/HOP4/model/model.py @@ -397,8 +397,15 @@ class SADLNet(nn.Sequential): model_for_export, self.get_example_inputs(patch_size, batch_size), filename, - input_names=["in"], - dynamic_axes={"in": {2: "h", 3: "w"}}, + input_names=["in0", "in1", "in2", "in3", "in4", "in5"], + dynamic_axes={ + "in0": {2: "h", 3: "w"}, + "in1": {2: "h", 3: "w"}, + "in2": {2: "h", 3: "w"}, + "in3": {2: "h", 3: "w"}, + "in4": {2: "h", 3: "w"}, + "in5": {2: "h", 3: "w"}, + }, opset_version=opset, **kwargs, )