Added support for Cuda GPU

iSiddharth20 · Dec 24, 2023 · 4c4428b · 4c4428b
1 parent 747e92c
commit 4c4428b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/Code/training.py b/Code/training.py
@@ -14,8 +14,10 @@
 # Define Training Class
 class Trainer():
     def __init__(self, model, loss_function, model_save_path):
-        # Define the model
-        self.model = model
+        # Define the device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        # Define the model and move it to the device
+        self.model = model.to(self.device)
         # Define the loss function
         self.loss_function = loss_function
         # Define the optimizer