[Python] How to find MSE when using a batch loader?

I'm working on a regression task using deep learning models. While calculating the MSE, I have divided by the length of the dataset. However, ChatGPT is suggesting me to divide it by the length of the loader. What is the correct way to calculate MSE in this case?

Also, I want to track the train_mse to see if the model is overfitting. While calculating the train_mse, should the model be in train or eval mode?

Code:

def train_and_evaluate(model, train_loader, val_loader, num_epochs=40, lr=1e-3, weight_decay=1e-5, patience=5):
model = model.to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

best_val_mse = float('inf')
epochs_no_improve = 0

for epoch in range(num_epochs):
model.train()
train_loss, train_preds, train_targets = 0, [], []
for features, targets in train_loader:
features, targets = features.to(device), targets.to(device)
preds = model(features)
loss = criterion(preds, targets)

optimizer.zero_grad()
loss.backward()
optimizer.step()

train_loss += loss.item()
train_preds.extend(preds.detach().cpu().numpy())
train_targets.extend(targets.cpu().numpy())

model.eval()
val_loss, val_preds, val_targets = 0, [], []
with torch.no_grad():
for features, targets in val_loader:
features, targets = features.to(device), targets.to(device)
preds = model(features)
loss = criterion(preds, targets)

val_loss += loss.item()
val_preds.extend(preds.detach().cpu().numpy())
val_targets.extend(targets.cpu().numpy())

train_mse = train_loss / len(train_dataset)
train_pc = safe_pearsonr(train_preds, train_targets)
val_mse = val_loss / len(val_dataset)
val_pc = safe_pearsonr(val_preds, val_targets)

Continue reading...