Incorrect scheduler stepping and label dtype handling in training loop

[SarthakJagota]

While reviewing the training pipeline, I noticed two small issues that may affect training stability:

1. The scheduler is stepped using scheduler.step(loss), but schedulers such as CosineAnnealingWarmRestarts do not require the loss value and expect a simple scheduler.step() call.
2. Labels are converted using labels.type(torch.LongTensor).to(device), which can create CPU tensors before moving to the target device. Using labels.long().to(device) is safer and avoids unnecessary tensor transfers.

I can submit a small fix addressing both points if this sounds reasonable.

[karthikduppala22-kar]

Hi! I have submitted a fix for this issue in PR #168.

The change updates scheduler stepping to scheduler.step() for compatibility with CosineAnnealingWarmRestarts and ensures correct label dtype handling.

Looking forward to feedback.

[likitha487]

Hello, I am interested in working on this issue.
I would like to contribute to this project as part of GSoC preparation.
Could you please assign this issue to me?

[likitha487]

Hi, I have submitted a PR to fix this issue.

PR: #173

Please let me know if any changes are required. Thank you!