Hi,
.register_buffer() is a function on nn.Modules, not autograd.Function.
These are two completely different constructs. And you won’t need any buffer for custom Function as it is not a class instance, you should just create the Tensor you need during the forward.