Fix scheduler stepping and label dtype handling in train.py and eval.py (#152)

DeepLense_Classification_Transformers_Archil_Srivastava/eval.py

@@ -47,9 +47,8 @@ def evaluate(model, data_loader, loss_fn, device):
     # Iterate over batches and accumulate metrics
     for batch_X, batch_y in data_loader:
         # Send data to device
-        batch_X, batch_y = batch_X.to(device, dtype=torch.float), batch_y.type(
-            torch.LongTensor
-        )
+        batch_X = batch_X.to(device, dtype=torch.float)
+        batch_y = batch_y.to(dtype=torch.long)
         logits.append(model(batch_X).cpu())  # Append the logits
         y.append(batch_y)  # Append the predictions
 

DeepLense_Classification_Transformers_Archil_Srivastava/train.py

@@ -44,17 +44,16 @@ 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
-    ).to(device)
+    images = images.to(device, dtype=torch.float)
+    labels = labels.to(device, dtype=torch.long)
     model.train()  # Set train mode
     optimizer.zero_grad()  # Reset gradients
     logits = model(images)  # Forward pass
     loss = criterion(logits, labels)  # Compute loss
     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