Ah yes, I read through his implementation. I think the main problem is that I wasn’t calculating the mask in the same latent space as the latents themselves, as @matdmiller mentioned here:
That is, my latents are 64x64x4 whereas I my mask is 512x512x3. So to apply the mask, I uncompressed my latents, applied the mask, and then recompressed the latents. Since the compression is lossy, I think that’s why the issue is occurring.
There are a few other differences from my implementation to the actual steps in the paper, but I don’t think they should make much of a difference.
I’ve decided to leave my implementation for now (have spent 2 weeks on it heh) and writing a post on my current implementation. Perhaps I’ll return to it later on in the course.