Skip to content

Commit

Permalink
allow auraflow latest to be torch compile compatible.
Browse files Browse the repository at this point in the history
  • Loading branch information
sayakpaul committed Jul 13, 2024
1 parent bbd2f9d commit 4faafa7
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def forward(self, latent):
)
latent = latent.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
latent = self.proj(latent)
return latent + self.pos_embed
return latent + self.pos_embed[:, : latent.size(1)]


# Taken from the original Aura flow inference code.
Expand Down

0 comments on commit 4faafa7

Please sign in to comment.