diff --git a/pyproject.toml b/pyproject.toml
index d61f841..5040633 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.0.1"
+version = "3.0.2"
 license = {file = "LICENSE"}
 
 [project.urls]
diff --git a/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py
index c7adbf5..a789243 100644
--- a/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py
+++ b/tests/pytorch_wrapper/pt_train_classification_conv_model_test2.py
@@ -67,7 +67,7 @@ class TestPtTrainConvModel2(unittest.TestCase):
                        use_gpu=True,
                        early_stopping=False,
                        early_stopping_rounds=1,
-                       output_best_val_model=False
+                       output_best_val_model=False,
                        val_loader=self.val_loader)[0]
         self.save_model.f(trained_model, f"dog_cat_{epochs}_epochs_conv.pt")