Tensor Size Mismatch

[akaul0]

I downloaded this repo and am trying to run it on my gpu but I get this error: The size of tensor a (32) must match the size of tensor b (64) at non-singleton dimension 3
It seems like the latent_prompt shape is this torch.Size([1, 2048, 64]) but it fails on the self.cfm.sample step because it expects that shape to be torch.Size([1, 2048, 32]) not 64. I am very confused and I don't see any one else that ran into this issue. Am I doing something wrong or missing something? Would appreciate some help thanks!