overfit

[yh-pengtu]

when I execute the overfit code as show in readme, it always output problems: flowmap.visualization.layout.hcat.
The problem arose whilst typechecking parameter 'images'.
Actual value: (f323,128,256, f321,3,128,128)
Expected type: typing.Iterable[Float[Tensor, 'channel _ _']].

[demplo]

For me, removing the Iterable type hinting from all the various cat functions solved this.

flowmap/flowmap/visualization/layout.py

Line 120 in 578a515

*images: Iterable[Float[Tensor, "channel _ _"]],

new type being *images: Float[Tensor, "channel _ _"],. Same for hcat and vcat.

Not sure this is the way it should be but jaxtyping stopped complaining.