fix branch input size for multi-channel inputs

Signed-off-by: nick7 <[email protected]>
boschresearch · Nov 28, 2023 · 8a90383 · 8a90383
1 parent ab2b264
commit 8a90383
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/torchphysics/models/deeponet/branchnets.py b/src/torchphysics/models/deeponet/branchnets.py
@@ -25,7 +25,7 @@ def __init__(self, function_space, discretization_sampler):
         super().__init__(function_space, output_space=None)
         self.output_neurons = 0
         self.discretization_sampler = discretization_sampler
-        self.input_dim = len(self.discretization_sampler)
+        self.input_dim = len(self.discretization_sampler) * function_space.output_space.dim
         self.current_out = torch.empty(0)
 
     def finalize(self, output_space, output_neurons):