Training stops trying to train on PointOdyssey

Hello! Thank you for your awesome work first of all. I am trying to train TAPTR on PointOdyssey, but I'm running into an issue during the first epoch - I believe even after running the first sample through the model. Here are the relevant lines from the log: <details><summary>Log</summary> ``` ... Start training Loss is nan, stopping training {'pt_full_cardinality_error_dn': tensor(0., device='cuda:0') 'pt_full_cardinality_error_dn_0': tensor(0., device='cuda:0') 'pt_full_cardinality_error_dn_1': tensor(0., device='cuda:0') 'pt_full_cardinality_error_dn_2': tensor(0., device='cuda:0') 'pt_full_cardinality_error_dn_3': tensor(0., device='cuda:0') 'pt_full_cardinality_error_dn_4': tensor(0., device='cuda:0') 'pt_full_loss_bbox': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_0': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_1': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_2': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_3': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_4': tensor(nan, device='cuda:0') 'pt_full_loss_bbox_dn': tensor(0., device='cuda:0') 'pt_full_loss_bbox_dn_0': tensor(0., device='cuda:0') 'pt_full_loss_bbox_dn_1': tensor(0., device='cuda:0') 'pt_full_loss_bbox_dn_2': tensor(0., device='cuda:0') 'pt_full_loss_bbox_dn_3': tensor(0., device='cuda:0') 'pt_full_loss_bbox_dn_4': tensor(0., device='cuda:0') 'pt_full_loss_ce': tensor(1.0215, device='cuda:0') 'pt_full_loss_ce_0': tensor(1.0736, device='cuda:0') 'pt_full_loss_ce_1': tensor(1.2466, device='cuda:0') 'pt_full_loss_ce_2': tensor(1.1136, device='cuda:0') 'pt_full_loss_ce_3': tensor(1.2177, device='cuda:0') 'pt_full_loss_ce_4': tensor(1.2346, device='cuda:0') 'pt_full_loss_ce_dn': tensor(0., device='cuda:0') 'pt_full_loss_ce_dn_0': tensor(0., device='cuda:0') 'pt_full_loss_ce_dn_1': tensor(0., device='cuda:0') 'pt_full_loss_ce_dn_2': tensor(0., device='cuda:0') 'pt_full_loss_ce_dn_3': tensor(0., device='cuda:0') 'pt_full_loss_ce_dn_4': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn_0': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn_1': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn_2': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn_3': tensor(0., device='cuda:0') 'pt_full_loss_giou_dn_4': tensor(0., device='cuda:0') 'pt_full_loss_hw': tensor(nan, device='cuda:0') 'pt_full_loss_hw_0': tensor(nan, device='cuda:0') 'pt_full_loss_hw_1': tensor(nan, device='cuda:0') 'pt_full_loss_hw_2': tensor(nan, device='cuda:0') 'pt_full_loss_hw_3': tensor(nan, device='cuda:0') 'pt_full_loss_hw_4': tensor(nan, device='cuda:0') 'pt_full_loss_hw_dn': tensor(0., device='cuda:0') 'pt_full_loss_hw_dn_0': tensor(0., device='cuda:0') 'pt_full_loss_hw_dn_1': tensor(0., device='cuda:0') 'pt_full_loss_hw_dn_2': tensor(0., device='cuda:0') 'pt_full_loss_hw_dn_3': tensor(0., device='cuda:0') 'pt_full_loss_hw_dn_4': tensor(0., device='cuda:0') 'pt_full_loss_xy': tensor(nan, device='cuda:0') 'pt_full_loss_xy_0': tensor(nan, device='cuda:0') 'pt_full_loss_xy_1': tensor(nan, device='cuda:0') 'pt_full_loss_xy_2': tensor(nan, device='cuda:0') 'pt_full_loss_xy_3': tensor(nan, device='cuda:0') 'pt_full_loss_xy_4': tensor(nan, device='cuda:0') 'pt_full_loss_xy_dn': tensor(0., device='cuda:0') 'pt_full_loss_xy_dn_0': tensor(0., device='cuda:0') 'pt_full_loss_xy_dn_1': tensor(0., device='cuda:0') 'pt_full_loss_xy_dn_2': tensor(0., device='cuda:0') 'pt_full_loss_xy_dn_3': tensor(0., device='cuda:0') 'pt_full_loss_xy_dn_4': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn_0': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn_1': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn_2': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn_3': tensor(0., device='cuda:0') 'pt_window_cardinality_error_dn_4': tensor(0., device='cuda:0') 'pt_window_loss_bbox': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_0': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_1': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_2': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_3': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_4': tensor(nan, device='cuda:0') 'pt_window_loss_bbox_dn': tensor(0., device='cuda:0') 'pt_window_loss_bbox_dn_0': tensor(0., device='cuda:0') 'pt_window_loss_bbox_dn_1': tensor(0., device='cuda:0') 'pt_window_loss_bbox_dn_2': tensor(0., device='cuda:0') 'pt_window_loss_bbox_dn_3': tensor(0., device='cuda:0') 'pt_window_loss_bbox_dn_4': tensor(0., device='cuda:0') 'pt_window_loss_ce': tensor(0.0751, device='cuda:0') 'pt_window_loss_ce_0': tensor(0.0818, device='cuda:0') 'pt_window_loss_ce_1': tensor(0.0988, device='cuda:0') 'pt_window_loss_ce_2': tensor(0.0852, device='cuda:0') 'pt_window_loss_ce_3': tensor(0.0929, device='cuda:0') 'pt_window_loss_ce_4': tensor(0.0972, device='cuda:0') 'pt_window_loss_ce_dn': tensor(0., device='cuda:0') 'pt_window_loss_ce_dn_0': tensor(0., device='cuda:0') 'pt_window_loss_ce_dn_1': tensor(0., device='cuda:0') 'pt_window_loss_ce_dn_2': tensor(0., device='cuda:0') 'pt_window_loss_ce_dn_3': tensor(0., device='cuda:0') 'pt_window_loss_ce_dn_4': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn_0': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn_1': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn_2': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn_3': tensor(0., device='cuda:0') 'pt_window_loss_giou_dn_4': tensor(0., device='cuda:0') 'pt_window_loss_hw': tensor(nan, device='cuda:0') 'pt_window_loss_hw_0': tensor(nan, device='cuda:0') 'pt_window_loss_hw_1': tensor(nan, device='cuda:0') 'pt_window_loss_hw_2': tensor(nan, device='cuda:0') 'pt_window_loss_hw_3': tensor(nan, device='cuda:0') 'pt_window_loss_hw_4': tensor(nan, device='cuda:0') 'pt_window_loss_hw_dn': tensor(0., device='cuda:0') 'pt_window_loss_hw_dn_0': tensor(0., device='cuda:0') 'pt_window_loss_hw_dn_1': tensor(0., device='cuda:0') 'pt_window_loss_hw_dn_2': tensor(0., device='cuda:0') 'pt_window_loss_hw_dn_3': tensor(0., device='cuda:0') 'pt_window_loss_hw_dn_4': tensor(0., device='cuda:0') 'pt_window_loss_xy': tensor(nan, device='cuda:0') 'pt_window_loss_xy_0': tensor(nan, device='cuda:0') 'pt_window_loss_xy_1': tensor(nan, device='cuda:0') 'pt_window_loss_xy_2': tensor(nan, device='cuda:0') 'pt_window_loss_xy_3': tensor(nan, device='cuda:0') 'pt_window_loss_xy_4': tensor(nan, device='cuda:0') 'pt_window_loss_xy_dn': tensor(0., device='cuda:0') 'pt_window_loss_xy_dn_0': tensor(0., device='cuda:0') 'pt_window_loss_xy_dn_1': tensor(0., device='cuda:0') 'pt_window_loss_xy_dn_2': tensor(0., device='cuda:0') 'pt_window_loss_xy_dn_3': tensor(0., device='cuda:0') 'pt_window_loss_xy_dn_4': tensor(0., device='cuda:0')} ``` </details> When I was preparing the PointOdyssey train videos, I followed the implementation of `__get_item__` in `PointTrackingDataset` in `kubric.py`, so that the format of samples and targets returned by my `PointOdysseyDataset` are the same as `PointTrackingDataset`. I'm fairly certain it should all match. Other than that, there are no other changes of significance, and the file `config/TAPTR.py` hasn't been changed in any way. Here's the command I used to launch the training (2 H100 80GB GPUs are used): ```bash python -m torch.distributed.launch --nproc_per_node=2 main.py \ -c config/TAPTR.py \ --dataset_file point_odyssey \ --data_path /path/to/pointodyssey \ --output_dir logs/train_taptr \ --num_workers 2 \ --options num_samples_per_video=56 num_queries_per_video=128 ``` `num_samples_per_video=56` and `num_queries_per_video=128` are set that way because the PointOdyssey training sequences I have are 56 frames long, each with 128 points. For this run, I only used 10 sequences just to check if everything goes through. Here are the raw requirements I installed in a Python 3.10 environment (that's what's available in the cluster) using CUDA 12.2.2: <details><summary>Requirements</summary> ```python torch==2.3.1 torchvision==0.18.1 numpy==1.26.4 tqdm opencv-python moviepy mediapy matplotlib gradio gradio-image-prompter timm scipy # MultiScaleDeformableAttention like in deformable DETR addict yapf==0.40.1 # https://github.com/open-mmlab/mmdetection/issues/10962 pycocotools termcolor albumentations tensorboard ``` </details> Yes, I am aware that my versions of most things are different, but I didn't manage to find any reason why they wouldn't work. Do you maybe have any ideas what I should check, or why this is happening? I would love to be able to run training in debug mode, but I don't think this is possible on an HPC cluster. Thank you in advance!