several errors when I train this model with re10k dataset

[dongshen6]

Hello,

Thank you for your excellent open-source work. However, while training with the re10k dataset, I encountered several errors that may have been caused by inconsistencies between the code and the dataset. I have made some modifications to my code and hope they can be helpful to other researchers.

1、The images parameter of cat,vcat,hcat expect an Iterable[Float[Tensor, 'channel _ _']], but the actual input is a tuple (f32[3,256,256], f32[3,256,256]). However, jaxtyping considers that an Iterable does not directly match a tuple (even though a tuple is a subclass of Iterable in Python, jaxtyping might enforce stricter type checks). So we can make the following modification:

2、The start parameter of the draw_lines function is expected to be one of the following types:
• float / int (a scalar)
• Iterable[Union[float, int]] (an iterable of numeric values)
• Shaped[Tensor, '3'] (a tensor of shape [3])
• Shaped[Tensor, 'batch 3'] (a tensor of shape [batch_size, 3])

However, the actual input is f32[24,2] (a tensor of shape [24, 2]), which does not meet the expected type. So we can make the following modification:

3、According to the definition of Vector, it allows the following types:

Vector = Union[
Real, # float or int (scalar)
Iterable[Real], # e.g., [0.1, 0.2, 0.3] or (0.1, 0.2, 0.3)
Shaped[Tensor, "3"], # A tensor of shape [3]
Shaped[Tensor, "batch 3"], # A tensor of shape [N, 3]
]

However, your inputs start and end are of type f32[24,2] (a tensor of shape [24, 2]), which does not match the existing definition of Vector because Vector requires the last dimension to be 3, while your data has 2.

Root Cause

The implementation logic of draw_lines is designed for 2D:
• sanitize_vector(start, 2, device) explicitly requires start to be 2D coordinates ([x, y]).
• However, the type definition of Vector enforces 3D ([x, y, z]), causing a conflict.

Issue

There is a mismatch between the type definition (Vector) and the actual requirement:
• Vector is designed for 3D data (e.g., RGB colors or 3D coordinates).
• But draw_lines expects 2D coordinates.