def training_step(self, batch, batch_idx):
(xb, yb), records = batch
preds = self(xb, yb)
loss = loss_fn(preds, yb)
self.log("train_loss", loss)
return loss
def validation_step(self, batch, batch_idx):
(xb, yb), records = batch
with torch.no_grad():
self.train()
train_preds = self(xb, yb)
loss = loss_fn(train_preds, yb)
self.eval()
raw_preds = self(xb)
preds = self.convert_raw_predictions(raw_preds=raw_preds)
self.accumulate_metrics(records=records, preds=preds)
self.log("val_loss", loss)