From c97e89d7e28654dd9fca5225c50023b40c570ba2 Mon Sep 17 00:00:00 2001 From: Shamsuddin Ahmed Date: Mon, 16 Sep 2024 12:08:13 +0600 Subject: [PATCH] set cuda device --- app/extractor.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/app/extractor.py b/app/extractor.py index 9a7df1d..1ab5f37 100644 --- a/app/extractor.py +++ b/app/extractor.py @@ -34,8 +34,12 @@ def __init__(self, config, video_path, frame_rate, output_dir, model_path, class self.image_processor = ImageProcessor(output_size=self.transformations.get('size', (640, 640))) # Set the device (CUDA or CPU) - self.device = 'cuda' if torch.cuda.is_available() else 'cpu' - + # Ensure CUDA is available + if torch.cuda.is_available(): + torch.cuda.set_device(0) # Sets GPU 0 (modify if using multiple GPUs) + self.device = 'cuda' + else: + self.device = 'cpu' # Only initialize SahiUtils if SAHI is enabled if sahi_config: self.sahi_utils = SahiUtils(self.config.debug, self.supported_classes_map,