diff --git a/docs/test_runner.py b/docs/test_runner.py
index cee079e2fa..29eb056551 100644
--- a/docs/test_runner.py
+++ b/docs/test_runner.py
@@ -1,18 +1,34 @@
-import os
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
 import subprocess
 
 import pytest
 
 # Collecting all the Python files in the scripts directory
-script_files = [f for f in os.listdir("test_scripts/") if f.endswith(".py")]
+test_scripts_dir = Path(__file__).parent / "test_scripts"
+script_files = [f for f in test_scripts_dir.glob("*.py")]
+
 
-# print(script_files)
 @pytest.mark.parametrize("script_file", script_files)
 def test_script_execution(script_file):
     """Test if a script runs without error."""
-    filepath = os.path.join("test_scripts/", script_file)
-    print(filepath)
-    result = subprocess.run(["python", filepath], capture_output=True, text=True)
+    print(f"Running {script_file}")
+    result = subprocess.run(["python", script_file], capture_output=True, text=True)
 
     # Check that the script executed successfully
     assert (
diff --git a/docs/test_scripts/test_custom_model_demo_1.py b/docs/test_scripts/test_custom_model_demo_1.py
index 8de8ba9cfe..e0337c8d30 100644
--- a/docs/test_scripts/test_custom_model_demo_1.py
+++ b/docs/test_scripts/test_custom_model_demo_1.py
@@ -70,11 +70,14 @@ class MdlsUNetMetaData(ModelMetaData):
 # [physicsnemo model]
 
 # [physicsnemo sym model]
-
-from typing import Dict, Optional
-
-from physicsnemo.sym.key import Key
-from physicsnemo.sym.models.arch import Arch
+try:
+    from physicsnemo.sym.key import Key
+    from physicsnemo.sym.models.arch import Arch
+except ImportError as e:
+    raise ImportError(
+        "This optional test depends on the `physicsnemo.sym` package.\n"
+        "If desired, you can install it with instructions from https://github.com/NVIDIA/physicsnemo-sym."
+    ) from e
 
 
 class MdlsSymUNet(Arch):
@@ -91,7 +94,7 @@ def __init__(
 
         self.mdls_model = MdlsUNet(in_channels, out_channels)  # MdlsUNet defined above
 
-    def forward(self, dict_tensor: Dict[str, torch.Tensor]):
+    def forward(self, dict_tensor: dict[str, torch.Tensor]):
         x = self.concat_input(
             dict_tensor,
             self.input_key_dict,