TrainingExtensions/torch/test/python/test_model_preparer.py
def train(model: torch.nn.Module, data_loader: DataLoader) -> torch.Tensor:
    """
    Helper function to train model given data loader
    :param model: torch model
    :param data_loader: torch data loader
    :return: total loss
    """
    total_loss = 0
    model.train()
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    for (data, labels) in data_loader:
        optimizer.zero_grad()
        predicted = model(data)
        loss = criterion(predicted, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss

    return total_loss

    Helper function to train model given data loader
    :param model: torch model
    :param data_loader: torch data loader
    :return: total loss
    
Helper function to train model given data loader
Helper function to train model given data loader
def train(model: torch.nn.Module, data_loader: DataLoader) -> torch.Tensor:
    total_loss = 0
    model.train()
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    for (data, labels) in data_loader:
        optimizer.zero_grad()
        predicted = model(data)
        loss = criterion(predicted, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss
    return total_loss
def train ( model : torch . nn . Module , data_loader : DataLoader ) -> torch . Tensor : total_loss = 0 model . train ( ) criterion = torch . nn . CrossEntropyLoss ( ) optimizer = torch . optim . SGD ( model . parameters ( ) , lr = 0.01 ) for ( data , labels ) in data_loader : optimizer . zero_grad ( ) predicted = model ( data ) loss = criterion ( predicted , labels ) loss . backward ( ) optimizer . step ( ) total_loss += loss return total_loss
Helper function to train model given data loader
Helper function to train model given data loader
"""
    Helper function to train model given data loader
    :param model: torch model
    :param data_loader: torch data loader
    :return: total loss
    """
None
model
torch.nn.Module
None
data_loader
DataLoader
torch data loader
torch data loader