Hardcode in `AttentionBlock`

[cuttle-fish-my]

Hi! Thanks for this fantastic work!

I am little confused about the forward function of AttentionBlock in unet.py. The corresponding codes are shown as blew:

 def forward(self, x):
        return checkpoint(self._forward, (x,), self.parameters(), True)

I just wonder why we should hardcode the flag parameter to True instead of usingself.use_checkpoints?

Thanks!

[wpMelene]

I think it is because this is a pixel-wise self attention implementation, which considered one pixel as a token. This implementation will be extremely memory comsuming when input resolution is relatively high (e.g. 256^2).