fix: indexing bugs from new shift functions

skhrg · Jan 28, 2024 · ba78497 · ba78497
1 parent 7a92115
commit ba78497
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/megham/transform.py b/megham/transform.py
@@ -154,8 +154,7 @@ def get_rigid(
 
     if row_basis:
         rot = rot.T
-    else:
-        shift = shift[..., np.newaxis]
+        shift = shift[:, 0]
 
     return rot, shift
 
@@ -235,8 +234,7 @@ def get_affine(
 
     if row_basis:
         affine = affine.T
-    else:
-        shift = shift[..., np.newaxis]
+        shift = shift[:, 0]
 
     return affine, shift