From 9d08ddf172e3f462a1252f3b97c0d7d83e4a88c0 Mon Sep 17 00:00:00 2001 From: Likitha Date: Mon, 9 Mar 2026 21:35:28 +0530 Subject: [PATCH] Fix scheduler stepping and label dtype handling in training loop --- .../train.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py index a5a6303..3961565 100644 --- a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py +++ b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py @@ -44,8 +44,7 @@ def train_step(model, images, labels, optimizer, scheduler, criterion, device="c Loss value from the forward pass """ # Send to device - images, labels = images.to(device, dtype=torch.float), labels.type( - torch.LongTensor + images, labels = images.to(device, dtype=torch.float), labels.long().to(device) ).to(device) model.train() # Set train mode optimizer.zero_grad() # Reset gradients @@ -54,7 +53,7 @@ def train_step(model, images, labels, optimizer, scheduler, criterion, device="c loss.backward() # Backward pass optimizer.step() # Optimize weights step if scheduler is not None: - scheduler.step(loss) # Modify learning rate if scheduler is set + scheduler.step() # Modify learning rate if scheduler is set return loss