Closed doctorpangloss closed 1 week ago
how do you know? in my testing torch.compile seems to unroll it into the equivalent vanilla operations and dropping einops did not make a different
that's good, I think this is more of an issue on Windows where torch.compile will not do that
https://github.com/replicate/cog-flux/blob/4c52f930e536bbe7d891c507d570667c803c0765/flux/modules/layers.py#L232
Replacing the einops
rearrange
with vanillatorch
will improve speed.