Fixed minor bug in data generator

weiya711 · Nov 5, 2024 · 044717c · 044717c
1 parent 87d8f2d
commit 044717c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/sam/onyx/generate_matrices.py b/sam/onyx/generate_matrices.py
@@ -88,8 +88,11 @@ def _create_fiber_tree(self):
         # self.fiber_tree = FiberTree(tensor=self.array if self.block_naive else self.array[...,
                                     # self.shape[-2]:self.block_size,:self.shape[-1]:self.block_size])
         self.fiber_tree = FiberTree(tensor=self.array)
-        self.tmp_fiber_tree = FiberTree(tensor=self.array[..., :self.shape[-2]:self.block_size,
-                                        :self.shape[-1]:self.block_size])
+        if len(self.shape) > 1:
+            self.tmp_fiber_tree = FiberTree(tensor=self.array[..., :self.shape[-2]:self.block_size,
+                                            :self.shape[-1]:self.block_size])
+        else:
+            self.tmp_fiber_tree = FiberTree(tensor=self.array)
 
     def transpose_tensor(self, axes=None):
         '''