Changed code to detach logits and labels to cpu before converting to …

…np array
advaithsrao · Nov 24, 2023 · c915715 · c915715
1 parent fd0e3a2
commit c915715
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/detector/modeler.py b/detector/modeler.py
@@ -1,5 +1,5 @@
 import os
-# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
 
 import shutil
 import pandas as pd
@@ -196,8 +196,8 @@ def train(
                 # Calculate the loss using the weighted loss function
                 loss = loss_function(logits, b_labels_one_hot)
                 total_eval_loss += loss.item()
-                logits = logits.detach().to(self.device).numpy()
-                label_ids = b_labels.to(self.device).numpy()
+                logits = logits.detach().cpu().numpy()
+                label_ids = b_labels.detach().cpu().numpy()
                 total_eval_accuracy += self.accuracy(logits, label_ids)
 
             avg_val_accuracy = total_eval_accuracy / len(validation_dataloader)
@@ -490,8 +490,8 @@ def train(
                 # Calculate the loss using the weighted loss function
                 loss = loss_function(logits, b_labels_one_hot)
                 total_eval_loss += loss.item()
-                logits = logits.detach().to(self.device).numpy()
-                label_ids = b_labels.to(self.device).numpy()
+                logits = logits.detach().cpu().numpy()
+                label_ids = b_labels.detach().cpu().numpy()
                 total_eval_accuracy += self.accuracy(logits, label_ids)
 
             avg_val_accuracy = total_eval_accuracy / len(validation_dataloader)