Hi,
You can find more details on how to implement custom functions here 5.
In particular, you don’t need the __init__ method.
And also, because you’re writing the custom Function, you only need to mark an output as non-differentiable if you don’t need gradient for it. Here you do, so you should remove the self.mark_non_differentiable(out) line.