Added automatic device detection for model loading in vulnscan

Resolved crash when loading PyTorch model on systems without CUDA by implementing auto-detection of device availability (CPU/GPU). torch.load now uses map_location to ensure compatibility across environments.

Signed-off-by: Shahm Najeeb <[email protected]>

Author: DefinetlyNotAI

CODE/vulnscan.py

@@ -55,7 +55,15 @@ def _load_model(self) -> None:
             elif self.model_path.endswith('.pth'):
                 with warnings.catch_warnings():
                     warnings.filterwarnings("ignore", category=FutureWarning)
-                    self.model = torch.load(self.model_path, weights_only=False)
+                self.model = torch.load(
+                    self.model_path,
+                    map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+                    weights_only=False
+                )
+                if not torch.cuda.is_available() and torch.version.cuda and torch.backends.cudnn.is_available():
+                    log.warning(
+                        "NVIDIA GPU detected but CUDA is not available. Check your PyTorch and CUDA installation to utilise as much power as possible.")
+                log.debug(f"Model using device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
             else:
                 raise ValueError("Unsupported model file format")