diff --git a/lietorch/models/ds_pde.py b/experiments/ds_pde.py
similarity index 100%
rename from lietorch/models/ds_pde.py
rename to experiments/ds_pde.py
diff --git a/experiments/run_ds_pde.py b/experiments/run_ds_pde.py
index 6ba5781eaa90fd5764646c80d5b6cfe10b0208cc..444851d7f0ca27f27b926e8ebff208420843095d 100644
--- a/experiments/run_ds_pde.py
+++ b/experiments/run_ds_pde.py
@@ -1,4 +1,4 @@
-from lietorch.models import ds_pde
+from . import ds_pde
 
 from dataclasses import asdict
 
@@ -105,11 +105,12 @@ if __name__ == "__main__":
     parser.add_argument(
         "--runname", type=str, default="", help="Name of the training run."
     )
+    parser.add_argument("--tag", type=str, default=None, help="wandb tag.")
     parser.add_argument(
-        "--tag", type=str, default=None, help="wandb tag."
-    )
-    parser.add_argument(
-        "--fixed_lift", type=eval, default=True, help="Lift with cakewavelets or trained wavelets."
+        "--fixed_lift",
+        type=eval,
+        default=True,
+        help="Lift with cakewavelets or trained wavelets.",
     )
 
     args = parser.parse_args()
diff --git a/lietorch/models/__init__.py b/lietorch/models/__init__.py
index e6f50fac93435d832f0c16cb2ac23f7e551bdd8c..ae25a5c550010d07ea6c779d45173ab21d2f571f 100644
--- a/lietorch/models/__init__.py
+++ b/lietorch/models/__init__.py
@@ -6,4 +6,3 @@ import lietorch.models.dca1
 import lietorch.models.drive
 import lietorch.models.rotnist
 import lietorch.models.fixed_lift
-import lietorch.models.ds_pde