Compute full training loss after every epoch

[fabian-sp]

I try to compute the loss and accuracy on the training set after each epoch. This is for vit_tiny_patch16_224.

Note: the train_metrics compute the loss over the course of one epoch, so its not the same as loss at end of epoch I guess.

I tried to simply add

train_metrics = validate(
                model,
                loader_train,
                train_loss_fn,
                args,
                amp_autocast=amp_autocast,
)

But this gives me an error:

RuntimeError: The expanded size of the tensor (512) must match the existing size (512000) at non-singleton dimension 1. Target sizes: [5, 512]. Tensor sizes: [1, 512000]

when computing the accuracy on the training set. I assume this is because the training set has some different structure than validation set.

Does someone have a suggestion how to do this?