def evaluate_loss( model: ModelInterface, batch: Sequence[Any], label: torch.Tensor ) -> torch.Tensor: # criterion = torch.nn.CrossEntropyLoss(weight=torch.tensor([1.0, 60.0])) criterion = torch.nn.CrossEntropyLoss() pred = model.forward(batch) loss = criterion(pred, label) return loss
def train_step( model: ModelInterface, # `torch.optim.optimizer.Optimizer` is ghost. # WHY DOES MYPY NOT RECOGNIZE `torch.optim.Optimizer`? optimizer: 'torch.optim.optimizer.Optimizer', batch: Sequence[Any], label: torch.Tensor) -> float: # criterion = torch.nn.CrossEntropyLoss(weight=torch.tensor([1.0, 60.0])) criterion = torch.nn.CrossEntropyLoss() optimizer.zero_grad() pred = model.forward(batch) loss = criterion(pred, label) loss.backward() optimizer.step() return loss.item()