Added Pt Train Model node to allow specifying a Loss node for training.

docs/reference/node_reference.md

@@ -225,6 +225,7 @@
 | [Pt Predict Regression Model](pt_predict_regression_model.md) | Performs inference on input data. |
 | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs.   |
 | [Pt Train Classification Model Lr](pt_train_classification_model_lr.md) | Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay.   |
+| [Pt Train Model](pt_train_model.md) | Trains a model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs.   |
 | [Pto Lr Scheduler Cosine Annealing](pto_lr_scheduler_cosine_annealing.md) | Creates a cosine annealing learning rate scheduler for an optimizer. |
 | [Pto Lr Scheduler Reduce On Plateau](pto_lr_scheduler_reduce_on_plateau.md) | Creates a reduce-on-plateau learning rate scheduler for an optimizer. |

docs/reference/pt_train_model.md

@@ -0,0 +1,34 @@
+# Pt Train Model
+Trains a model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.  
+
+early_stopping_rounds: Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped.
+output_best_val_model: Whether to outputs the model with the best val loss or the model from the last epoch.
+
+Functionality-wise, this node is identical with PtTrainClassificationModelLr except that you can specify a custom loss function in this model where PtTrainClassificationModelLr uses nn.CrossEntropyLoss().
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+| train_loader | Ptdataloader |
+| optimizer | Ptoptimizer |
+| loss_function | Ptloss |
+| epochs | Int |
+| use_gpu | Boolean |
+| early_stopping | Boolean |
+| early_stopping_rounds | Int |
+| output_best_val_model | Boolean |
+| scheduler | Ptlrscheduler |
+| val_loader | Ptdataloader |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+| Tensor |
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.