From 3af7b1427475158c814f05b71cbda255e4c1c2fe Mon Sep 17 00:00:00 2001
From: Eric Bezzam <ebezzam@gmail.com>
Date: Thu, 8 Aug 2024 02:31:20 -0700
Subject: [PATCH] Fix for random shifts. (#139)

---
 lensless/recon/utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/lensless/recon/utils.py b/lensless/recon/utils.py
index ca25a704..65c6e238 100644
--- a/lensless/recon/utils.py
+++ b/lensless/recon/utils.py
@@ -474,6 +474,7 @@ def __init__(
         clip_grad=1.0,
         unrolled_output_factor=False,
         random_rotate=False,
+        random_shift=False,
         pre_proc_aux=False,
         extra_eval_sets=None,
         use_wandb=False,
@@ -554,7 +555,6 @@ def __init__(
         post_process_unfreeze : int, optional
             Epoch at which to unfreeze post process component. Default is None.
 
-
         """
         global print
 
@@ -612,6 +612,9 @@ def __init__(
             self.train_multimask = train_dataset.multimask
         self.train_random_flip = train_dataset.random_flip
         self.random_rotate = random_rotate
+        self.random_shift = random_shift
+        if self.random_shift:
+            raise NotImplementedError("Random shift not implemented yet.")
 
         # check if Subset and if simulating dataset
         self.simulated_dataset_trainable_mask = False