update code

deepxde/data/pde_operator.py

@@ -263,28 +263,34 @@ def _losses(self, outputs, loss_fn, inputs, model, num_func, aux=None):
             # Use stack instead of as_tensor to keep the gradients.
             losses = [bkd.reduce_mean(bkd.stack(loss, 0)) for loss in losses]
         elif config.autodiff == "forward":  # forward mode AD
-            outputs_shape = bkd.shape(outputs)
-            shape0, shape1 = outputs_shape[0], outputs_shape[1]
-            shape2 = model.net.num_outputs
+            shape0, shape1 = bkd.shape(outputs)[:2]
 
             def forward_call(trunk_input):
                 output = aux[0]((inputs[0], trunk_input))
-                return bkd.reshape(output, (shape0 * shape1, shape2))
+                return bkd.reshape(output, (shape0 * shape1, model.net.num_outputs))
 
             f = []
             if self.pde.pde is not None:
                 # Each f has the shape (N1, N2)
                 f = self.pde.pde(
                     inputs[1],
-                    (bkd.reshape(outputs, (shape0 * shape1, shape2)), forward_call),
-                    bkd.reshape(model.net.auxiliary_vars, (shape0 * shape1, shape2)),
+                    (
+                        bkd.reshape(outputs, (shape0 * shape1, model.net.num_outputs)),
+                        forward_call,
+                    ),
+                    bkd.reshape(
+                        model.net.auxiliary_vars,
+                        (shape0 * shape1, model.net.num_outputs),
+                    ),
                 )
                 if not isinstance(f, (list, tuple)):
                     f = [f]
             f = (
                 [bkd.reshape(fi, (shape0, shape1)) for fi in f]
                 if model.net.num_outputs == 1
-                else [bkd.reshape(fi, (shape0, shape1, shape2)) for fi in f]
+                else [
+                    bkd.reshape(fi, (shape0, shape1, model.net.num_outputs)) for fi in f
+                ]
             )
 
             # Each error has the shape (N1, ~N2)