torch
pytorch-lightning
torchmetrics
