fix cp import logic

TorchDSP · Feb 6, 2023 · 87d357e · 87d357e
1 parent 8d7ea1a
commit 87d357e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/torchsig/datasets/synthetic.py b/torchsig/datasets/synthetic.py
@@ -11,6 +11,7 @@
 try:
     import cusignal
     import cupy as xp
+    import cupy as cp
     CUSIGNAL = True
     CUPY = True
 except ImportError:  
@@ -738,7 +739,7 @@ def _generate_samples(self, item: Tuple) -> np.ndarray:
 
         symbols = const_oversampled[symbol_nums]
         symbols_repeat = xp.repeat(symbols, samples_per_symbol_recalculated)
-        symbols_repeat = xp.insert(symbols_repeat,0,0) # start at zero phase
+        symbols_repeat = np.insert(symbols_repeat,0,0) # start at zero phase
 
         filtered = symbols_repeat
         if "g" in const_name:

diff --git a/torchsig/datasets/wideband.py b/torchsig/datasets/wideband.py
@@ -12,6 +12,7 @@
 try:
     import cusignal
     import cupy as xp
+    import cupy as cp
     CUSIGNAL = True
     CUPY = True
 except ImportError: