From 6a4810f79e87a3dedc568a023a76005fb3428ccb Mon Sep 17 00:00:00 2001
From: Manuja5d <165181756+Manuja5d@users.noreply.github.com>
Date: Sun, 19 Jan 2025 20:43:11 -0500
Subject: [PATCH] Update main.py

---
 main.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/main.py b/main.py
index b9eba81..f206759 100644
--- a/main.py
+++ b/main.py
@@ -10,6 +10,12 @@
 if __name__ == '__main__':
     # Check if CUDA is available and print the result
     print("CUDA Available: ", torch.cuda.is_available())
+    logging.info("CUDA Available: %s", torch.cuda.is_available())
+    device = torch.device("cuda" if (cfg.use_cuda and torch.cuda.is_available()) else "cpu")
+    if cfg.use_cuda and not torch.cuda.is_available():
+        logging.warning("CUDA is enabled in the configuration, but no CUDA devices are available. Falling back to CPU.")
+    else:
+        logging.info("Using device: %s", device)
     # Set the device to CUDA if available and configured to use, otherwise use CPU
     device = torch.device("cuda" if (cfg.use_cuda and torch.cuda.is_available()) else "cpu")