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()
