diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py b/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
index a3f12b3..5913e13 100644
--- a/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
+++ b/DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
@@ -56,9 +56,10 @@ def evaluate(model, data_loader, loss_fn, device):
     # Concatenate all results
     logits, y = torch.cat(logits), torch.cat(y)
     loss.append(loss_fn(logits, y))
-    accuracy.append(accuracy_fn(logits, y, num_classes=NUM_CLASSES))
-    class_auroc.append(auroc_fn(logits, y, num_classes=NUM_CLASSES, average=None))
-    macro_auroc.append(auroc_fn(logits, y, num_classes=NUM_CLASSES, average="macro"))
+    accuracy.append(accuracy_fn(logits, y, task="multiclass", num_classes=NUM_CLASSES))
+    class_auroc.append(auroc_fn(logits, y, task="multiclass", num_classes=NUM_CLASSES, average=None))
+    micro_auroc.append(auroc_fn(logits, y, task="multiclass", num_classes=NUM_CLASSES, average="micro"))
+    macro_auroc.append(auroc_fn(logits, y, task="multiclass", num_classes=NUM_CLASSES, average="macro"))
 
     result = {
         "ground_truth": y,
@@ -93,7 +94,7 @@ def evaluate(model, data_loader, loss_fn, device):
 
     # Start wandb run
     with wandb.init(
-        entity="_archil", project=run_config.project, id=run_config.runid, resume="must"
+        project=run_config.project, id=run_config.runid, resume="must"
     ):
         # Get best device on machine
         device = get_device(run_config.device)
@@ -169,7 +170,7 @@ def evaluate(model, data_loader, loss_fn, device):
         roc_auc = dict()
         for idx, cls in enumerate(LABELS):
             class_truth = (metrics["ground_truth"].numpy() == idx).astype(int)
-            class_pred = torch.nn.functional.softmax(metrics["logits"]).numpy()[
+            class_pred = torch.nn.functional.softmax(metrics["logits"], dim=-1).numpy()[
                 ..., idx
             ]
             fpr[idx], tpr[idx], _ = roc_curve(class_truth, class_pred)
diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
index a5a6303..795f7d7 100644
--- a/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
+++ b/DeepLense_Classification_Transformers_Archil_Srivastava/train.py
@@ -234,7 +234,6 @@ def train(
 
     # Start wandb run
     with wandb.init(
-        entity="_archil",
         project=run_config.project,
         config=run_config,
         group=group,
@@ -319,7 +318,7 @@ def train(
         # Scheduler
         if run_config.decay_lr:
             scheduler = CosineAnnealingWarmRestarts(
-                optimizer, T_0=15, T_mult=1, eta_min=1e-6, verbose=True
+                optimizer, T_0=15, T_mult=1, eta_min=1e-6
             )
         else:
             scheduler = None
diff --git a/DeepLense_Classification_Transformers_Archil_Srivastava/utils.py b/DeepLense_Classification_Transformers_Archil_Srivastava/utils.py
index 759d3fe..aeb97eb 100644
--- a/DeepLense_Classification_Transformers_Archil_Srivastava/utils.py
+++ b/DeepLense_Classification_Transformers_Archil_Srivastava/utils.py
@@ -29,7 +29,7 @@ def get_device(device):
         return xm.xla_device()
     if (device == "cuda" or device == "best") and torch.cuda.is_available():
         return "cuda"
-    if (device == "mps" or device == "best") and torch.has_mps:
+    if (device == "mps" or device == "best") and torch.backends.mps.is_built():
         return "mps"
     if device == "cpu" or device == "best":
         return "cpu"