After running train nothing is happening. Stuck on Running on Cuda.

Hello,

I'm trying to recreate the experiments and after running train nothing is happening:

!python experiments/train.py --dataset_dir ../datasets --train_dataset synthetic --N 120000 --B 20 --total_steps 601 --encoder_spec ALLMINILM --key_embd_src key --use_data_aug --log_to_file --verbose

Output:

wandb: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.
wandb: (1) Create a W&B account
wandb: (2) Use an existing W&B account
wandb: (3) Don't visualize my results
wandb: Enter your choice: 3
wandb: You chose "Don't visualize my results"
wandb: Tracking run with wandb version 0.19.9
wandb: W&B syncing is set to offline in this directory.
wandb: Run wandb online or set WANDB_MODE=online to enable cloud syncing.
[10:01:24] INFO Running on cuda
---- Just stops here ---

The log file shows:

2025-04-10 10:00:08,854 - training - INFO - Running on cuda
2025-04-10 10:00:08,857 - training - INFO - 🚨 Started training 🚨
2025-04-10 10:00:08,859 - training - INFO - 💽 Saving to output💽
2025-04-10 10:00:08,860 - training - INFO - Using pre-computed ALLMINILM embedding
--- Stops here ---

Using Google Colab with T4.

Created embeddings using all-MiniLM-L6-v2.

Could you help me with this ?

Reproduce the command in experiments/train.py on the reported Colab T4 setup, then trace the execution after the “Using pre-computed ALLMINILM embedding” log entry. Check whether training advances or blocks at that point and document the cause. Done means training starts progressing or the failure has a clear, reproducible diagnosis.