diff --git a/aiida_cp2k/utils/input_generator.py b/aiida_cp2k/utils/input_generator.py
index dc17e6f..4806bf8 100644
--- a/aiida_cp2k/utils/input_generator.py
+++ b/aiida_cp2k/utils/input_generator.py
@@ -213,13 +213,14 @@ def add_ext_restart_section(input_dict, first_snapshot=None):
     if first_snapshot is not None:
         params["MOTION"]["MD"]["REFTRAJ"]["FIRST_SNAPSHOT"] = first_snapshot
     # overwrite the complete EXT_RESTART section if present
-    params["EXT_RESTART"] = {"RESTART_FILE_NAME": "./parent_calc/aiida-1.restart"}
-    if params["GLOBAL"]["RUN_TYPE"] == "MD":
-        params["EXT_RESTART"]["RESTART_DEFAULT"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_COUNTERS"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_POS"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_VEL"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_CELL"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_THERMOSTAT"] = ".TRUE."
-        params["EXT_RESTART"]["RESTART_CONSTRAINT"] = ".FALSE."
+    params["EXT_RESTART"] = {
+        "RESTART_FILE_NAME": "./parent_calc/aiida-1.restart",
+        "RESTART_DEFAULT": ".TRUE.",
+        "RESTART_COUNTERS": ".TRUE.",
+        "RESTART_POS": ".TRUE.",
+        "RESTART_VEL": ".TRUE.",
+        "RESTART_CELL": ".TRUE.",
+        "RESTART_THERMOSTAT": ".TRUE.",
+        "RESTART_CONSTRAINT": ".FALSE.",
+    }
     return Dict(params)
diff --git a/aiida_cp2k/workchains/base.py b/aiida_cp2k/workchains/base.py
index 4bd420f..b36468c 100644
--- a/aiida_cp2k/workchains/base.py
+++ b/aiida_cp2k/workchains/base.py
@@ -95,7 +95,7 @@ def restart_incomplete_calculation(self, calc):
                 first_snapshot = int(params['MOTION']['MD']['REFTRAJ']['FIRST_SNAPSHOT']) + calc.outputs.output_trajectory.get_shape('positions')[0]
             except KeyError:
                 pass
-            params = add_ext_restart_section(params,first_snapshot=first_snapshot)
+            params = add_ext_restart_section(params, first_snapshot=first_snapshot)
 
         self.ctx.inputs.parameters = params  # params (new or old ones) that include the necessary restart information.
         self.report(
diff --git a/examples/workchains/example_base_md_reftraj_restart.py b/examples/workchains/example_base_md_reftraj_restart.py
index b0dd07b..f62a7e0 100644
--- a/examples/workchains/example_base_md_reftraj_restart.py
+++ b/examples/workchains/example_base_md_reftraj_restart.py
@@ -17,7 +17,6 @@
 from aiida.engine import run
 from aiida.orm import Dict, SinglefileData, load_code
 from aiida.plugins import DataFactory, WorkflowFactory
-from ase import Atoms
 
 Cp2kBaseWorkChain = WorkflowFactory("cp2k.base")
 StructureData = DataFactory("core.structure")