Merge pull request #116 from ttngu207/dev_use_cuda_flag

feat: add `use_cuda` for caiman
datajoint · Jul 11, 2024 · 86f1f85 · 86f1f85
2 parents e025b29 + f3aea53
commit 86f1f85
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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"