Fix scheduler stepping and label dtype handling in training loop (#152)

[karthikduppala22-kar]

Fix scheduler stepping and label dtype handling in the training loop.

Changes:

• Replaced scheduler.step(loss) with scheduler.step() for compatibility with CosineAnnealingWarmRestarts.
• Ensured labels are converted to torch.long before computing loss.
• Returned scalar loss using loss.item() for correct logging.

Fixes #152.