Hi Zimo!



 ZimoNitrome:

some of the color neurons output 1.0, and torch.logit(1.0) gives inf.


Instead of applying sigmoid() to all of your neurons and then applying
logit() to your color neurons to undo the sigmoid(), apply sigmoid()
only to your non-color neurons so that the color neurons never get
transformed.
(Applying sigmoid() to a large-enough (non-inf) value will map it
to 1.0, causing some loss of information, and you can’t undo the
numerical “damage” by applying logit(), even though mathematically
logit() is the inverse function of sigmoid().)
As an aside, you should consider whether you really want to apply
sigmoid() to any of your neurons.  What subsequent processing
do you do to your non-color neurons, and could you perform that
processing in “log space,” that is, working with the logits that are
output by your Linear head?  Doing so might prove more numerically
stable.
Best.
K. Frank