 you’ll see the timestep weights are calculated and added as noise:   Copied weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps) If you want to learn more about how the training loop works, check out the Understanding pipeline