diff --git a/more_math/LatentMathNode.py b/more_math/LatentMathNode.py
index 2dddb80..08327a1 100644
--- a/more_math/LatentMathNode.py
+++ b/more_math/LatentMathNode.py
@@ -154,7 +154,7 @@ class LatentMathNode(MathNodeBase):
             merged_result = eval_single_tensor(merged_a, merged_b, merged_c, merged_d)
 
             split_results = list(merged_result.split(sizes, dim=0))
-            out_samples = comfy.nested_tensor.nested_tensor.NestedTensor(split_results)
+            out_samples = comfy.nested_tensor.NestedTensor(split_results)
 
             return ({"samples": out_samples},)
 
diff --git a/pytest.ini b/pytest.ini
index b1d9480..f43d8e4 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,2 +1,2 @@
 [pytest]
-pythonpath = . d:/stability/Data/Packages/ComfyUI/custom_nodes/more_math d:/stability/Data/Packages/ComfyUI
+pythonpath = . ../..
diff --git a/tests/conftest.py b/tests/conftest.py
index 72f3a07..5d75628 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -13,6 +13,28 @@ for p in [_project_root, _comfy_root]:
         # also make it visible to subprocesses that inspect PYTHONPATH
         os.environ["PYTHONPATH"] = p + os.pathsep + os.environ.get("PYTHONPATH", "")
 
+# Special case: 'more_math' root folder has an __init__.py which conflicts with the 'more_math' subfolder.
+# We want 'import more_math' to resolve to the subfolder package.
+if "more_math" not in sys.modules or "custom_nodes" in (getattr(sys.modules["more_math"], "__file__", "") or ""):
+    import importlib.util
+    _subfolder = os.path.join(_project_root, "more_math")
+    _spec = importlib.util.spec_from_file_location("more_math", os.path.join(_subfolder, "__init__.py"))
+    if _spec:
+        _m = importlib.util.module_from_spec(_spec)
+        sys.modules["more_math"] = _m
+        _spec.loader.exec_module(_m)
+
+# Redirect comfy_api.Model to comfy_api.latest.io.Model if needed
+# as requested by user to support the 'latest' API position.
+try:
+    import comfy_api
+    import comfy_api.latest as latest
+    if not hasattr(comfy_api, "Model") and hasattr(latest, "io") and hasattr(latest.io, "Model"):
+        comfy_api.Model = latest.io.Model
+        comfy_api.io = latest.io
+except (ImportError, AttributeError):
+    pass
+
 import pytest
 
 
diff --git a/tests/test_model_math.py b/tests/test_model_math.py
index 52a3c29..6847e4a 100644
--- a/tests/test_model_math.py
+++ b/tests/test_model_math.py
@@ -22,7 +22,7 @@ from more_math.ModelMathNode import ModelMathNode
 class MockModelPatcher:
     def __init__(self, state_dict):
         self.model = comfy_api.Model()
-        self.model.state_dict = state_dict
+        self.model.state_dict = lambda: state_dict
         self.patches = {}
 
     def clone(self):
@@ -163,7 +163,7 @@ class MockPatcherContainer:
 
 
 def test_clip_math():
-    from more_math.ModelMathNode import CLIPMathNode
+    from more_math.ClipMathNode import CLIPMathNode
 
     sd_a = {"text_model.encoder.layers.0.mlp.fc1.weight": torch.tensor([1.0])}
     clip_a = MockPatcherContainer(sd_a)
@@ -178,7 +178,7 @@ def test_clip_math():
 
 
 def test_vae_math():
-    from more_math.ModelMathNode import VAEMathNode
+    from more_math.VaeMathNode import VAEMathNode
 
     sd_a = {"decoder.conv_in.weight": torch.tensor([1.0])}
     # VAE acts like PatcherContainer but no clone() on the VAE itself usually.
diff --git a/tests/test_more_math.py b/tests/test_more_math.py
index 769ff50..aa975bf 100644
--- a/tests/test_more_math.py
+++ b/tests/test_more_math.py
@@ -260,7 +260,7 @@ def test_image_swap():
     node = ImageMathNode()
     img_red = torch.tensor([1.0, 0.0, 0.0]).view(1, 1, 1, 3)
     img_blue = torch.tensor([0.0, 0.0, 1.0]).view(1, 1, 1, 3)
-    res_swap = node.execute("swap(a, 1, 0, 2)", a=img_red)[0]
+    res_swap = node.execute("swap(a, 3, 0, 2)", a=img_red)[0]
     assert torch.allclose(res_swap, img_blue)