From f3aea5354df761d80f520e75b916e5c819352883 Mon Sep 17 00:00:00 2001
From: Thinh Nguyen <thinh@datajoint.com>
Date: Sat, 6 Jul 2024 08:22:44 -0500
Subject: [PATCH] feat: add `use_cuda` for caiman

---
 element_interface/run_caiman.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/element_interface/run_caiman.py b/element_interface/run_caiman.py
index d77a9e2..51fb044 100644
--- a/element_interface/run_caiman.py
+++ b/element_interface/run_caiman.py
@@ -8,6 +8,13 @@
 except:  # noqa E722
     pass  # TODO: remove bare except
 
+try:
+    import torch
+    cuda_is_available = torch.cuda.is_available()
+except:
+    cuda_is_available = False
+    pass
+
 import caiman as cm
 from caiman.source_extraction.cnmf import params as params
 from caiman.source_extraction.cnmf.cnmf import CNMF
@@ -36,6 +43,8 @@ def run_caiman(
     parameters["fnames"] = file_paths
     parameters["fr"] = sampling_rate
 
+    parameters["use_cuda"] = cuda_is_available  # Use CUDA if available
+
     if "indices" in parameters:
         indices = parameters.pop(
             "indices"