def train_epoch()

in lab/wind_turbine.py [0:0]

• 17 lines of code
• 5 McCabe index (conditional complexity)

def train_epoch(optimizer, criterion, epoch, model, train_dataloader, test_dataloader):
    train_loss = 0.0    
    test_loss = 0.0    
    model.train()
    for x_train, y_train in train_dataloader:
        # clearing the Gradients of the model parameters
        optimizer.zero_grad()
        # prediction for training and validation set        
        output_train = model(x_train)        
        loss_train = criterion(output_train, y_train)
                
        # computing the updated weights of all the model parameters
        # statistics
        train_loss += loss_train.item()
        loss_train.backward()
        optimizer.step()        
    model.eval()
    for x_test, y_test in test_dataloader:            
        output_test = model(x_test.float())
        loss_test = criterion(output_test, y_test)
        # statistics
        test_loss += loss_test.item()                
        
    return train_loss, test_loss