I have a 2-layers fully connected network. I would like to convert the output of the first layer to binary. This means that I would like to have a binary-step activation function in the forward paths and Relu activation function in the backward pass. How can I implement this? Any idea would be appreciated.
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 64)
        self.fc2 = nn.Linear(64, 10)
        nn.init.xavier_normal_(self.fc1.weight)
        nn.init.xavier_normal_(self.fc2.weight)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        x = F.log_softmax(x, dim=1)
        return x

net = Net()