Question: torch.max term used in `softmax_kernel` about performer-pytorch HOT 4 CLOSED

@lucidrains Thanks for your response!

Sorry to bother you, but I have another question about the implementation in reversible.py:

        with torch.enable_grad():
            x2.requires_grad = True
            fx2 = self.f(x2, set_rng=True, **f_args)
            torch.autograd.backward(fx2, dx1, retain_graph=True)

Is param retain_graph=True necessary here?

from performer-pytorch.

@ClawangTU yes it is, it is necessary for the scenario where both the encoder and decoder are reversible

from performer-pytorch.

@ClawangTU I don't really know either, I transcribed this faithfully from their repository https://github.com/google-research/google-research/blob/master/performer/fast_attention/jax/fast_attention.py#L100-L107

from performer-pytorch.

@lucidrains I get it! Thanks a lot for your answer!

from performer-pytorch.