diff --git a/comfy_Extras/nodes/nodes_train.py b/comfy_extras/nodes/nodes_train.py
similarity index 99%
rename from comfy_Extras/nodes/nodes_train.py
rename to comfy_extras/nodes/nodes_train.py
index a3951ac9b..10a70c559 100644
--- a/comfy_Extras/nodes/nodes_train.py
+++ b/comfy_extras/nodes/nodes_train.py
@@ -20,7 +20,7 @@ from comfy import node_helpers
 from comfy.cli_args import args
 from comfy.comfy_types.node_typing import IO
 from comfy.weight_adapter import adapters
-from .nodes import nodes_custom_sampler
+from . import nodes_custom_sampler
 
 
 class TrainSampler(comfy.samplers.Sampler):
diff --git a/tests/unit/comfy_extras_test/image_stitch_test.py b/tests/unit/comfy_extras_test/image_stitch_test.py
index d845c0a01..eb64516b5 100644
--- a/tests/unit/comfy_extras_test/image_stitch_test.py
+++ b/tests/unit/comfy_extras_test/image_stitch_test.py
@@ -1,6 +1,5 @@
 import pytest
 import torch
-from unittest.mock import MagicMock
 
 from comfy_extras.nodes.nodes_images import ImageStitch