I see your point! So, I guess the easiest way is to write custom layers and define the forward() functions to include the sign() before returning the output. So, for example, instead of using nn.Conv2d, you can define your own MyConv2d so that it applies the sign() function in its forward. Then, change all the layesr built using nn.Conv2d to MyConv2d. Similarly, you can define a MyLinear class that does the FC layer but in addition it applies the sign() function.