disable numba imports

requirements.txt

@@ -1,5 +1,6 @@
 interactive-pipe>=0.7.8
 torch>=2.0.0
 tqdm
+numba
 
 

scripts/save_deadleaves.py

@@ -9,12 +9,16 @@ from pathlib import Path
 from time import perf_counter
 import matplotlib.pyplot as plt
 from typing import Tuple
-
+import logging
 import numpy as np
 import torch
 import torch.nn.functional as F
 from torch.utils.data import Dataset
-from numba import cuda
+try:
+    from numba import cuda
+except ImportError:
+    logging.warning("Numba not installed, GPU acceleration will not be available")
+    cuda = None
 from tqdm import tqdm
 import argparse
 from rstor.synthetic_data.dead_leaves_gpu import gpu_dead_leaves_chart

src/rstor/synthetic_data/dead_leaves_gpu.py

@@ -3,7 +3,11 @@ from rstor.properties import SAMPLER_UNIFORM
 from typing import Tuple, Optional
 from rstor.synthetic_data.dead_leaves_cpu import define_dead_leaves_chart
 import numpy as np
-from numba import cuda
+try:
+    from numba import cuda
+except ImportError:
+    logging.warning("Numba not installed, GPU acceleration will not be available")
+    cuda = None
 import math
 
 

src/rstor/utils.py

@@ -1,5 +1,8 @@
 import numpy as np
-import numba
+try:
+    import numba
+except ImportError:
+    numba = None
 import torch
 
 THREADS_PER_BLOCK = 32  # 32 or 16