Training with Kitti dataset

Hi I am a newbie in this field and trying to train SiLK with Kitti odometry dataset.
Thanks for sharing your code!

# error i got
```
File "/usr/local/lib/python3.8/dist-packages/jax/_src/interpreters/pxla.py", line 1349, in __call__
results = self.xla_executable.execute_sharded(input_bufs)
│ │ └ [Array([[[-11.48745345, 23.01827276],
│ │ [-10.5068542 , 22.79328857],
│ │ [ -9.52790215, 22.56868231],
│ │ .....
│ └
└

jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Executable expected shape f64[1,21316,2]{2,1,0} for argument 0 but got incompatible shape f64[1,21316,2]{2,0,1}
2024-03-05 07:08:52.388 | ERROR | silk.cli:main:116 - run failed, `*.log` file might be found in : .

```

What does the {2,1,0} or {2,0,1} mean?
Is f64[1,21316,2] a feature vector?

I tried to put print(~~.shape) in your code but I can't even find it on terminal..

I don't have validation set like explained below.
Is this because of dataset setting?

Help me pls.

# dataset setting

Changed odometry dataset into torchvision.datasets.Kitti format like below cause you used torchvision dataset loader for other datasets.
Modified
As you see torchvision.datasets.Kitti has training, testing folder, not validation folder.

# environment

I use docker on ubuntu20.04
Docker base image: nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
**with no conda**

Other packages are installed with required version but...
Some packges below are different from requirements.
I struggled a lot setting environment and it was the first running environment for me.

jaxlib-0.4.13+cuda11.cudnn86
torch==2.2.0+cu118
torchvision==0.17.0+cu118

Start at silk.cli:main:116 and reproduce the training command with the reported Kitti-derived dataset and environment versions. Trace the dataset and training inputs until the reported JAX shape mismatch appears; done means identifying whether the dataset configuration or environment causes the failure and documenting a verified way to run training.

Written by the indexing model from the issue text.