Bug in ide.py when handling a list of losses

[isaingit]

Dear developers,

I found a bug in the code of ide.py, because it triggered an error like "list object is not callable" when an argument of type "loss = ['MSE', 'MSE'] is passed to the compile method.

Find below a proposal to overcome the issue:

class IDE(PDE):
     
    (...)

    def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):           
        bcs_start = np.cumsum([0] + self.num_bcs)
        int_mat = self.get_int_matrix(True)
        f = self.pde(inputs, outputs, int_mat)
        if not isinstance(f, (list, tuple)):
            f = [f]
        f = [fi[bcs_start[-1] :] for fi in f]
        
        if not isinstance(loss_fn, (list, tuple)):
            loss_fn = [loss_fn] * (len(f) + len(self.bcs))
        elif len(loss_fn) != len(f) + len(self.bcs):
            raise ValueError(
                "There are {} errors, but only {} losses.".format(
                    len(f) + len(self.bcs), len(loss_fn)
                )
            )
            
        # losses = [loss_fn(bkd.zeros_like(fi), fi) for fi in f]
        losses = [loss_fn[i](bkd.zeros_like(fi), fi) for i, fi in enumerate(f)]
        
        for i, bc in enumerate(self.bcs):
            beg, end = bcs_start[i], bcs_start[i + 1]
            error = bc.error(self.train_x, inputs, outputs, beg, end)
            # losses.append(loss_fn(bkd.zeros_like(error), error))
            losses.append(loss_fn[len(f) + i](bkd.zeros_like(error), error))
        return losses

    def losses_test(self, targets, outputs, loss_fn, inputs, model, aux=None):
        int_mat = self.get_int_matrix(False)
        f = self.pde(inputs, outputs, int_mat)
        if not isinstance(f, (list, tuple)):
            f = [f]
        # return [
        #     loss_fn(bkd.zeros_like(fi), fi) for fi in f
        # ] + [bkd.as_tensor(0, dtype=config.real(bkd.lib)) for _ in self.bcs]
        return [
            loss_fn[i](bkd.zeros_like(fi), fi) for i, fi in enumerate(f)
        ] + [bkd.as_tensor(0, dtype=config.real(bkd.lib)) for _ in self.bcs]

Would that be OK? Thanks in advance for the attention :)