lucidrains / meshgpt-pytorch

Implementation of MeshGPT, SOTA Mesh generation using Attention, in Pytorch
MIT License
700 stars 57 forks source link

`face_mask` rearrange in Autoencoder Forward Pass #40

Closed qixuema closed 8 months ago

qixuema commented 8 months ago

Hi Phil,

Happy New Year! 🎉🎉🎉

I encountered a minor issue in the code when I set return_recon_faces to True. During the execution of the autoencoder's forward method, I noticed that face_mask is rearranged.

face_mask = rearrange(face_mask, 'b nf -> b nf 1 1')
recon_faces = recon_faces.masked_fill(~face_mask, float('nan'))

However, after applying masked_fill to recon_faces, it seems that the shape of face_mask isn't restored. This is causing an error when trying to perform a repeat operation on face_mask later in the process.

face_mask = repeat(face_mask, 'b nf -> b (nf r)', r = 9)

Perhaps the shape of face_mask could be restored right after the masked_fill operation.

Best regards, Xueqi Ma

lucidrains commented 8 months ago

hey happy new years! yea i'll get back to this in jan

in the spirit of open source, always welcoming PRs

qixuema commented 8 months ago

Thank you for your prompt reply. I hope your holiday season is filled with joy and relaxation.

In the meantime, I'll consider working on a Pull Request to address this issue. Once I have something ready, I'll submit a PR for your review.