fixed small bug in Crossbar.py

The code would have called memtorch_cuda_bindings even if the version was CPU.
coreylammie · Feb 7, 2022 · de8dc2d · de8dc2d
1 parent 34a7862
commit de8dc2d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/memtorch/bh/crossbar/Crossbar.py b/memtorch/bh/crossbar/Crossbar.py
@@ -242,13 +242,17 @@ def write_conductance_matrix(
             if (
                 self.use_bindings
                 and type(self.devices.any()) in CUDA_passive_implemented
+                and "cpu" not in memtorch.__version__
             ):
                 device_matrix = build_g_tensor(self)
+                device_matrix_aug = device_matrix
+                conductance_matrix_aug = conductance_matrix
                 if (
                     len(device_matrix.shape) == 2
                 ):  # To ensure compatibility with CUDA code
                     device_matrix_aug = device_matrix[:, :, None]
                     conductance_matrix_aug = conductance_matrix[:, :, None]
+
                 new_matrix = memtorch_cuda_bindings.simulate_passive(
                     conductance_matrix_aug,
                     device_matrix_aug,