diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py b/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
index a3f12b32..78d92ac7 100644
--- a/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
+++ b/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
 
diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
index a5a6303c..e839ea1d 100644
--- a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
+++ b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
@@ -44,9 +44,8 @@ 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
@@ -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