Is your feature request related to a problem? Please describe.
If we have class API for metrics, it makes easier to set all training components before training loop.
Get 3 feedback from Andriy:

1. Add class API for the metrics functions.
2. Change torch.tensor(float(“nan”)).to(y_o.float()) to
   torch.tensor(float("nan"), device=y_o.device)
3. Add dim= to several PyTorch calls.