diff --git a/align.py b/align.py
index 2be1c85cf25e76526554ef7f02af5ce3085d421e..7963c3a407f45728134479d557f0bf728229f9f9 100644
--- a/align.py
+++ b/align.py
@@ -118,9 +118,9 @@ def align2D(
         
     data = stack_dask.compute()
     if rotation_data_deg != 0:
-        data = rotate(data, rotation_data_deg)
+        data = rotate(input=data, angle=rotation_data_deg, axes=(1,2))
         print(f'Rotated data {rotation_data_deg} deg')
-        
+
     aligned, tvec = register.align_stack(
         data, 
         path_to_save=None,