def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
"""
Actual metric computation
Args:
pred: predicted labels
target: ground truth labels
Return:
A Tensor with the mae loss.
"""
return mae(pred, target, self.reduction)
def __step(self, batch, batch_idx, prefix: str):
x = batch.x
edge_index = batch.edge_index
mini_batch = batch.batch
edge_attr = batch.edge_attr
y_hat = self.forward(x,
edge_index,
edge_attr,
mini_batch).squeeze(-1)
loss = pl_regression.mse(y_hat, batch.y)
rmse = torch.sqrt(loss)
mae = pl_regression.mae(y_hat, batch.y)
predictions = y_hat.detach().cpu().numpy()
r2 = r2_score(predictions, batch.y.cpu())
return {
f'{prefix}_loss': loss,
f'{prefix}_size': len(y_hat),
f'{prefix}_rmse': rmse,
f'{prefix}_mae': mae,
f'{prefix}_r2': r2,
}