Solved by crowsonkb in post #2 
              
              
                You can do projected gradient descent by enforcing your constraint after each optimizer step. An example training loop would be: 
    opt = optim.SGD(model.parameters(), lr=0.1)
    for i in range(1000):
        out = model(inputs)
        loss = loss_fn(out, labels)
        print(i, loss.item())
  …