ENH: compute loss for each batch

[Xylambda]

The actual version of the trainer computes the loss of a batch as the loss of the last computation of that batch.

See train and val.

It would be better to compute the mean of all losses of the batch. Then, the mean would be taken to get the average loss of a batch.

A possible code to do that:

    def _train(self, loader):
        self.model.train()
        _losses = []

        for features, labels in loader:
            # move to device
            features, labels = self._to_device(features, labels, self.device)
            
            # forward pass
            out = self.model(features)
            
            # loss
            loss = self._compute_loss(out, labels)
            
            # remove gradient from previous passes
            self.optimizer.zero_grad()
            
            # backprop
            loss.backward()
            
            # parameters update
            self.optimizer.step()

            _losses.append(loss.item())
            
        return np.array(_losses).mean()

[Xylambda]

Done in version 2.0.0