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

docs/model_training.md

@@ -95,3 +95,7 @@ Below is an example workflow for classifying the classic Iris dataset:
 `examples/workflows/iris_classification_train_eval.json`.
 
 According to [the UCI repository](https://archive.ics.uci.edu/dataset/53/iris), the neural network baseline accuracy for this dataset ranges from 92.105% to 100%. Using this workflow, you can expect around 96% accuracy, occasionally reaching 100%, which aligns with the baseline.
+
+# Specifying a loss function for training
+If you want to specify a loss function for training instead of using the default one in the trainer, use the `Pt Train Model` node. Create a loss function node (e.g., `Pt Cross Entropy Loss`) and connect it to the trainer node.
+Refer to `examples/workflows/train_model_with_custom_loss_example.json` for an example. (This workflow is intended to demonstrate node connections and is not optimized for high CIFAR10 accuracy. It runs for only 10 epochs to allow a quick check and reaches 88.9% accuracy.) You should adjust the hyperparameters according to your specific goals.

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.