train_val.py

@baker.command
def run(batch_size=8, epochs=5, device='cuda'):
  train_loader, val_loader = get_data_loaders(batch_size)
  model = NoiseClassifier()
  model = model.to(device)
  model.eval()
  
  optimizer = Adam(model.parameters())
  criterion = nn.BCELoss()
  
  trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
  
  val_metrics = {"accuracy": Accuracy(), "loss": Loss(criterion)}
  evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)
  
  # ...
  # Here we are attaching loggers
  # Look into the repository for full code
  # ...
  
  trainer.run(train_loader, max_epochs=epochs)