torch.bernoulli derivative is zero (due to discrete output)
your autograd.Function doesn’t make sense, because your “function”'s derivative doesn’t depend on its input. you don’t handle the problem of discreteness, and can as well supply any “external” binary tensor to the network with the same effect (no backprop beyond it).
RelaxedBernoulli is a workaround, that works by allowing continuous outputs (samples) during training. It is not an easy way, but you can read about it - [1611.00712] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables