Incorrect metrics reported at the end of validation epoch

In our example Hyrax models, often the validate_batch methods look something like this:

    def validate_batch(self, batch):
        _, labels = batch
        outputs = self(batch)
        loss = self.criterion(outputs, labels)
        return {"loss": loss.item()}

The loss value is just the loss calculated for the current batch. However, we advertise this as the loss at the end of validation, which a user could very reasonably assume means the total loss for the validation dataset.

We should actually be accumulating the loss over all the batches and then dividing by the number of batches (or number of batches and batch size if the final batch is not a complete batch)

Hyrax also exposes a validate_post_epoch method hook, where the correct, averaged metrics could be reported. e.g.:

def validate_batch(self, batch):
    _, labels = batch
    outputs = self(batch)
    loss = self.criterion(outputs, labels)
    self.total_val_loss += loss.item()
    self.num_batches += 1
    return {"loss": loss.item()}

def validate_post_epoch(self):
    avg_loss = self.total_val_loss / self.num_batches
    self.total_val_loss = 0
    self.num_batches = 0
    return {"avg_loss": avg_loss}

Start by finding the example Hyrax models that implement validate_batch and validate_post_epoch, then trace how validation metrics are collected and reported at epoch end. Ensure validation loss is averaged across the validation batches, including the final incomplete batch as appropriate, and verify that the reported post-epoch metrics match that average.