Fix scheduler stepping and label dtype handling in training loop

[likitha487]

Fixes issue #152

Changes made:

1. Replaced scheduler.step(loss) with scheduler.step() since some schedulers like CosineAnnealingWarmRestarts do not require the loss value.
2. Updated label conversion from labels.type(torch.LongTensor).to(device) to labels.long().to(device) to avoid unnecessary CPU tensor creation.

This improves training stability and prevents unnecessary tensor transfers.