diff --git a/suite2p/registration/nonrigid.py b/suite2p/registration/nonrigid.py
index e1e935d7..f371a60b 100644
--- a/suite2p/registration/nonrigid.py
+++ b/suite2p/registration/nonrigid.py
@@ -461,4 +461,4 @@ def transform_data(data, nblocks, xblock, yblock, ymax1, xmax1,
                              mode="bilinear", padding_mode="border", align_corners=True)
         
     
-    return fr_shift.squeeze().short()#.cpu().numpy()
+    return fr_shift.squeeze().to(dtype=data.dtype)