I want to do something like :
from torch.autograd import Function
class BinaryLayer(Function):
    def forward(self, input):
        return (input > .5).float()

    def backward(self, grad_output):
        return grad_output

Where the forward would be a binary activation and the backward would be linear.
I get the following error:
RuntimeError: could not compute gradients for some functions
when I use said function.
How can I get it to work?
Thanks in advance,
Dan