You can do this as follows:
X1, X2 = some tensors
loss_weighting = some constant to weight the different losses

X2.requires_grad_()

task_loss = NN(X1, X2)

gX2 = torch.autograd.grad(task_loss, X2, create_graph=True)[0]
penalty_loss = torch.max(-gX2, 0).sum() # other things than sum could be used here...

loss = task_loss + loss_weighting * penalty_loss

optimizer.zero_grad()
loss.backward()
optimizer.step()