def __init__(self, model: torch.nn.Module, loss: Loss) -> None:
self.model = model
self.loss = loss
self.criterion = loss._create_pytorch_loss()
def __init__(self, model: TorchModel, loss: Loss) -> None:
self.model = model
self.loss = loss # not used
self.criterion = loss._create_pytorch_loss()