Solved by KFrank in post #2 
              
              
                Hi Metro! 

I don’t understand your use case, but, in general, you can certainly 
penalize outputs that are not close to 0 or 1. 
Let x be a single output.  You could then try things like: 
penalty_0_1 = (x * (1.0 - x))**2
penalty_0_1 = torch.abs (x * (1.0 - x))
penalty_0_1 = ((x * (1.0 - x))**2)**a…