Model.train() and model.eval() when performing training and evaluation simultaneously

I recently learned that there are evaluation and training modes in the model. If I do training and evaluation at the same time to check the overtitting, where do I set model.eval() and model.train()?

Are the below codes correct?

# Train the model
oneEpochLossList_train = []
for i, batch in enumerate(train_loader):
    inputs, labels = batch

    # Move tensors to the configured device
    inputs = inputs.to(device)
    labels = labels.to(device)

    # Forward pass
    model.train()
    y_pred = model(inputs)
    loss = criterion(y_pred, labels)

    oneEpochLossList_train.append(loss.item())

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Test the model with current parameters
oneEpochLossList_test = []
for j, batch in enumerate(test_loader):
    inputs, labels = batch

    # Move tensors to the configured device
    inputs = inputs.to(device)
    labels = labels.to(device)

    # Forward pass
    model.eval()
    y_pred = model(inputs)
    loss = criterion(y_pred, labels)

    oneEpochLossList_test.append(loss.item())

And I want to know what the model.eval() and model.train() are exactly.
Thanks.

Your code works properly, but you can just put model.train() and model.eval() right before the corresponding for loop for training and validating, not do it for every loop.

# Train the model
oneEpochLossList_train = []
model.train()
for i, batch in enumerate(train_loader):
   ...

# Test the model with current parameters
oneEpochLossList_test = []
model.eval()
for j, batch in enumerate(test_loader):
    ...

1 Like

Thanks!
I edited it as you taught me.