Hi,
You should change your function train. If you have an issue doing this, please share your train function, and we can adapt it to do evaluation after few batches, in all cases I think you train function look like
def train(scheduler, optimizer):
    for x, y in train_loader:
        x = x.to(device)
        y = y.to(device)
        
       optimizer.zero_grad()
       outs = model(x)
       loss = criterion(outs, y)
       
       loss.backward()

       optimizer.step()
       scheduler.step()
  

You can update it and have something like
def train(scheduler, optimizer):
    for batch_idx, (x, y) in enumerate(train_loader):
        x = x.to(device)
        y = y.to(device)
        
       optimizer.zero_grad()
       outs = model(x)
       loss = criterion(outs, y)
       
       loss.backward()

       optimizer.step()
       scheduler.step()

       if batch_idx % log_freq == batch_idx -1:
           valid_loss = evaluate() 
           print(f"Batch {batch_idx} / {len(train_loader)} | Valid Loss = {valid_loss}")
           if valid_loss < best_valid_loss:
               best_valid_loss = valid_loss
               torch.save(model.state_dict(), model_file)