incerto.utils.train_epoch

Contents

incerto.utils.train_epoch#

incerto.utils.train_epoch(model, train_loader, criterion, optimizer, device, desc='Training')[source]#

Train model for one epoch.

Parameters:

  • model (Module) – Model to train

  • train_loader (DataLoader) – Training data loader

  • criterion (Module) – Loss function

  • optimizer (Optimizer) – Optimizer

  • device (device) – Device to train on

  • desc (str) – Description for progress bar

Return type:

dict

Returns:

Dictionary with ‘loss’ and ‘accuracy’

Contents