Add samples at all good, less budget

Change the file name

Changes in Dit block

Complete the script with wandb logging

Add the training script for dit