APE-T feature channels not reduced after forward()?

[lcxrocks]

First, thank the authors for their valuable contribution to the community!

In the paper, the learned residual is added to the refined training images cache F', and the affinity matrix is computed between f' and F', as described in equation (13) and beyond.

However, I noticed a small discrepancy between the paper and the code. In the code, the selected Q channels are only modified, not extracted as described in the paper.

Here's the relevant code:

APE/utils.py

Line 132 in 70460de

new_cache_keys[:, self.indices] = new_cache_keys[:, self.indices] + res_keys

Could this be due to a misunderstanding on my part, or is there an intentional reason for this difference? Thank you for your time in advance!