Update linear_transform.py

LiamBindle · May 12, 2022 · 1e537c5 · 1e537c5
1 parent e7affa4
commit 1e537c5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sparselt/linear_transform.py b/sparselt/linear_transform.py
@@ -23,11 +23,13 @@ def __init__(self, weights, row_ind, col_ind,
 
         self._input_core_dims = tuple(input_transform_dims[0])
         self._input_core_shape = tuple(input_transform_dims[1])
+        input_size = np.product(self._input_core_shape)
 
         self._output_core_dims = tuple(output_transform_dims[0])
         self._output_core_shape = tuple(output_transform_dims[1])
+        output_size = np.product(self._output_core_shape)
 
-        self._matrix = scipy.sparse.csr_matrix((weights, (row_ind, col_ind)))
+        self._matrix = scipy.sparse.csr_matrix((weights, (row_ind, col_ind)), shape=(output_size, input_size))
         self._order = order
 
         self._vfunc = self._create_vfunc()