def penalty()

in code/colored_mnist/main.py [0:0]

• 5 lines of code
• 1 McCabe index (conditional complexity)

  def penalty(logits, y):
    scale = torch.tensor(1.).cuda().requires_grad_()
    loss = mean_nll(logits * scale, y)
    grad = autograd.grad(loss, [scale], create_graph=True)[0]
    return torch.sum(grad**2)