Added simple optimizer and single layer linear model as well as a workflow to show case those nodes.

__init__.py

@@ -20,6 +20,7 @@ from .modules.pytorch_wrapper.pt_cos import PtCos
 from .modules.pytorch_wrapper.pt_cosh import PtCosh
 from .modules.pytorch_wrapper.pt_crop import PtCrop
 from .modules.pytorch_wrapper.pt_data_loader import PtDataLoader
+from .modules.pytorch_wrapper.pt_data_loader_from_tensors import PtDataLoaderFromTensors
 from .modules.pytorch_wrapper.pt_div import PtDiv
 from .modules.pytorch_wrapper.pt_einsum import PtEinsum
 from .modules.pytorch_wrapper.pt_eq import PtEq
@@ -64,6 +65,7 @@ from .modules.pytorch_wrapper.pt_pad import PtPad
 from .modules.pytorch_wrapper.pt_permute import PtPermute
 from .modules.pytorch_wrapper.pt_pow import PtPow
 from .modules.pytorch_wrapper.pt_predict_classification_model import PtPredictClassificationModel
+from .modules.pytorch_wrapper.pt_predict_regression_model import PtPredictRegressionModel
 from .modules.pytorch_wrapper.pt_prod import PtProd
 from .modules.pytorch_wrapper.pt_rand import PtRand
 from .modules.pytorch_wrapper.pt_rand_int import PtRandInt
@@ -94,18 +96,21 @@ from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy
 from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors
 from .modules.pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel
 from .modules.pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr
+from .modules.pytorch_wrapper.pt_train_regression_model import PtTrainRegressionModel
 from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze
 from .modules.pytorch_wrapper.pt_var import PtVar
 from .modules.pytorch_wrapper.pt_view import PtView
 from .modules.pytorch_wrapper.pt_where import PtWhere
 from .modules.pytorch_wrapper.pt_zeros import PtZeros
 from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
+from .modules.pytorch_wrapper.ptn_linear import PtnLinear
 from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel
 from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel
 from .modules.pytorch_wrapper.pto_adam import PtoAdam
 from .modules.pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing
 from .modules.pytorch_wrapper.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau
 from .modules.pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep
+from .modules.pytorch_wrapper.pto_simple import PtoSimple
 from .modules.pytorch_wrapper.ptv_dataset import PtvDataset
 from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen
 from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader
@@ -142,6 +147,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtCosh": PtCosh,
     "PtCrop": PtCrop,
     "PtDataLoader": PtDataLoader,
+    "PtDataLoaderFromTensors": PtDataLoaderFromTensors,
     "PtDiv": PtDiv,
     "PtEinsum": PtEinsum,
     "PtEq": PtEq,
@@ -186,6 +192,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtPermute": PtPermute,
     "PtPow": PtPow,
     "PtPredictClassificationModel": PtPredictClassificationModel,
+    "PtPredictRegressionModel": PtPredictRegressionModel,
     "PtProd": PtProd,
     "PtRand": PtRand,
     "PtRandInt": PtRandInt,
@@ -216,18 +223,21 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtToRgbTensors": PtToRgbTensors,
     "PtTrainClassificationModel": PtTrainClassificationModel,
     "PtTrainClassificationModelLr": PtTrainClassificationModelLr,
+    "PtTrainRegressionModel": PtTrainRegressionModel,
     "PtUnsqueeze": PtUnsqueeze,
     "PtVar": PtVar,
     "PtView": PtView,
     "PtWhere": PtWhere,
     "PtZeros": PtZeros,
     "PtnConvModel": PtnConvModel,
+    "PtnLinear": PtnLinear,
     "PtnLinearModel": PtnLinearModel,
     "PtnResnetModel": PtnResnetModel,
     "PtoAdam": PtoAdam,
     "PtoLrSchedulerCosineAnnealing": PtoLrSchedulerCosineAnnealing,
     "PtoLrSchedulerReduceOnPlateau": PtoLrSchedulerReduceOnPlateau,
     "PtoLrSchedulerStep": PtoLrSchedulerStep,
+    "PtoSimple": PtoSimple,
     "PtvDataset": PtvDataset,
     "PtvDatasetLen": PtvDatasetLen,
     "PtvDatasetLoader": PtvDatasetLoader,
@@ -264,6 +274,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtCosh": "Pt Cosh",
     "PtCrop": "Pt Crop",
     "PtDataLoader": "Pt Data Loader",
+    "PtDataLoaderFromTensors": "Pt Data Loader From Tensors",
     "PtDiv": "Pt Div",
     "PtEinsum": "Pt Einsum",
     "PtEq": "Pt Eq",
@@ -308,6 +319,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtPermute": "Pt Permute",
     "PtPow": "Pt Pow",
     "PtPredictClassificationModel": "Pt Predict Classification Model",
+    "PtPredictRegressionModel": "Pt Predict Regression Model",
     "PtProd": "Pt Prod",
     "PtRand": "Pt Rand",
     "PtRandInt": "Pt Rand Int",
@@ -338,18 +350,21 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtToRgbTensors": "Pt To Rgb Tensors",
     "PtTrainClassificationModel": "Pt Train Classification Model",
     "PtTrainClassificationModelLr": "Pt Train Classification Model Lr",
+    "PtTrainRegressionModel": "Pt Train Regression Model",
     "PtUnsqueeze": "Pt Unsqueeze",
     "PtVar": "Pt Var",
     "PtView": "Pt View",
     "PtWhere": "Pt Where",
     "PtZeros": "Pt Zeros",
     "PtnConvModel": "Ptn Conv Model",
+    "PtnLinear": "Ptn Linear",
     "PtnLinearModel": "Ptn Linear Model",
     "PtnResnetModel": "Ptn Resnet Model",
     "PtoAdam": "Pto Adam",
     "PtoLrSchedulerCosineAnnealing": "Pto Lr Scheduler Cosine Annealing",
     "PtoLrSchedulerReduceOnPlateau": "Pto Lr Scheduler Reduce On Plateau",
     "PtoLrSchedulerStep": "Pto Lr Scheduler Step",
+    "PtoSimple": "Pto Simple",
     "PtvDataset": "Ptv Dataset",
     "PtvDatasetLen": "Ptv Dataset Len",
     "PtvDatasetLoader": "Ptv Dataset Loader",

docs/model_training.md

@@ -3,10 +3,13 @@
 You can train a PyTorch model without any coding.
 Training is as simple as a single click!  
 Currently, classification models are supported for following architecture (node names are in parentheses):
-* Linear (Ptn Linear Model)
+* Linear for classificication (Ptn Linear Model)
 * Convolutional Network (Ptn Conv Model)
 * ResNet (Ptn Resnet Model)
 
+In addition, a single layer Linear model is supported as an example to illustrage the use (see the bottom of this page on how to use it):
+* Linear(Ptn Linear)
+
 You can adjust the model configuration including output dimension on the UI.
 
 The system is designed to be modular. For example, to switch from a linear model to a convolutional model, you only need to replace the model node (see screenshot below).
@@ -69,3 +72,19 @@ This train_v8 model achieved 93.12% accuracy & train_v11_2nd_run model achieved
 
 **ResNet model**
 ![Eval](images/resnet_eval.png)
+
+# Using the Single-Layer Linear Model
+
+A single-layer linear model is included to illustrate how to connect various nodes.  
+A workflow containing this model and related nodes is in the following file:  
+`examples/workflows/simple_linear_regression_train.json`.
+
+This workflow illustrates:
+- How to generate linear regression test data
+- How to plot the data
+- How to train a model by connecting it to an optimizer
+- How to make predictions using a trained model
+
+If you run training, you will notice that the loss remains in the 20s–30s instead of approaching 0, as you might expect.
+
+Do you know why this happens?

docs/reference/node_reference.md

@@ -132,12 +132,16 @@
 | Node | Description |
 | --- | --- |
 | [Pt Data Loader](pt_data_loader.md) | Loads data from a dataset node and creates a PyTorch DataLoader.   |
+| [Pt Data Loader From Tensors](pt_data_loader_from_tensors.md) | Creates a Torchvision Dataloader from a pair of tensors. |
 | [Pt Evaluate Classification Model](pt_evaluate_classification_model.md) | Performs inference on test data and computes evaluation metrics. |
 | [Pt Load Model](pt_load_model.md) | A wrapper class for saving a PyTorch model. |
 | [Pt Predict Classification Model](pt_predict_classification_model.md) | Performs inference on input data. |
+| [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 Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs.   |
 | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers.   |
+| [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer.   |
 | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers.   |
 | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers.   |
 | [Pto Adam](pto_adam.md) | Instantiates the Adam optimizer. |

docs/reference/pt_data_loader_from_tensors.md

@@ -0,0 +1,28 @@
+# Pt Data Loader From Tensors
+Creates a Torchvision Dataloader from a pair of tensors.
+
+Parameters  
+
+**x**  
+  Specify a tensor name that contains x.
+  
+**y**  
+  Specify a tensor name that contains y.
+
+## Input
+| Name | Data type |
+|---|---|
+| x | Tensor |
+| y | Tensor |
+| batch_size | Int |
+| shuffle | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptdataloader |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/pt_predict_regression_model.md

@@ -0,0 +1,14 @@
+# Pt Predict Regression Model
+Performs inference on input data.
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+| inputs | Tensor |
+| use_gpu | Boolean |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/pt_train_regression_model.md

@@ -0,0 +1,30 @@
+# Pt Train Regression Model
+Trains a regression model using a given dataset, optimizer, and number of epochs.  
+
+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.
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+| train_loader | Ptdataloader |
+| optimizer | Ptoptimizer |
+| epochs | Int |
+| use_gpu | Boolean |
+| early_stopping | Boolean |
+| early_stopping_rounds | Int |
+| output_best_val_model | Boolean |
+| val_loader | Ptdataloader |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+| Tensor |
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_linear.md

@@ -0,0 +1,24 @@
+# Ptn Linear
+A linear model consisting of a single dense layer.  
+
+### Fields:  
+- `in_features`: The number of input features.
+- `out_features`: The number of output features.
+- `bias`: Use bias or not.
+
+## Input
+| Name | Data type |
+|---|---|
+| in_features | Int |
+| out_features | Int |
+| bias | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/pt_data_loader_from_tensors.py

@@ -0,0 +1,67 @@
+from typing import Any, Dict, Tuple
+import torch
+from torch.utils.data import TensorDataset, DataLoader
+
+
+class PtDataLoaderFromTensors:
+    """
+    Creates a Torchvision Dataloader from a pair of tensors.
+    
+    Parameters  
+    
+    **x**  
+      Specify a tensor name that contains x.
+      
+    **y**  
+      Specify a tensor name that contains y.
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the required input types.  
+        
+        Returns:  
+            Dict[str, Any]: A dictionary of required input types.  
+        """
+        return {
+            "required": {
+                "x": ("TENSOR", {}),
+                "y": ("TENSOR", {}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "shuffle": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES: tuple = ("PTDATALOADER",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self, 
+          x: torch.Tensor,
+          y: torch.Tensor,
+          batch_size: int,
+          shuffle: bool) -> Tuple:
+        """
+        Creates a Torchvision Dataloader from a pair of tensors.
+        
+        Args:  
+            x (torch.Tensor): A tensor containing x.
+            y (torch.Tensor): A tensor containing y.
+            batch_size (int): The number of samples per batch.
+            shuffle (bool): If True, shuffles the dataset before each epoch.
+
+        Returns:  
+            Tuple: A tuple containing the dataset instance.  
+        """
+        
+        with torch.inference_mode(False):
+            dataset = TensorDataset(x, y)
+            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
+            return (dataloader,)

modules/pytorch_wrapper/pt_predict_classification_model.py

@@ -38,7 +38,7 @@ class PtPredictClassificationModel:
           class_id_to_name_map: Dict,
           use_gpu: bool) -> Tuple[str, int, float]:
         """
-        Performs inference on test data and computes evaluation metrics.
+        Performs inference on test data.
 
         Args:
             model (torch.nn.Module): The PyTorch model to evaluate.

modules/pytorch_wrapper/pt_predict_regression_model.py

@@ -0,0 +1,61 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class PtPredictRegressionModel:
+    """
+    Performs inference on input data.
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {}),
+                "inputs": ("TENSOR", {}),
+                "use_gpu": ("BOOLEAN", {"default":False})
+            }
+        }
+
+    RETURN_NAMES: Tuple[str] = ("y_hat",)
+    RETURN_TYPES: Tuple[str] = ("TENSOR",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    def f(self, model: nn.Module,
+          inputs: torch.Tensor,
+          use_gpu: bool) -> Tuple[torch.Tensor]:
+        """
+        Performs inference on test data.
+
+        Args:
+            model (torch.nn.Module): The PyTorch model to evaluate.
+            inputs (torch.Tensor): Input tensor.
+            use_gpu (bool): Whether to use GPU for inference.
+
+        Returns:
+            Tuple[torch.Tensor]: A tuple containing the prediction.
+        """
+
+        with torch.inference_mode():
+            if use_gpu:
+                model.to("cuda")
+            model.train(False)  # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
+            x = inputs
+            if x.dim() == 1:  # Add axis for col
+                x = torch.unsqueeze(x, -1)
+            if use_gpu:
+               x = x.to("cuda")
+            outputs = model(x)
+            if use_gpu:
+                outputs = outputs.to("cpu")
+        return (outputs,)

modules/pytorch_wrapper/pt_train_regression_model.py

@@ -0,0 +1,167 @@
+import os
+import sys
+from typing import Any, Dict, Optional, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class PtTrainRegressionModel:
+    """
+    Trains a regression model using a given dataset, optimizer, and number of epochs.  
+
+    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.
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types required for training the model.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {}),
+                "train_loader": ("PTDATALOADER", {}),
+                "optimizer": ("PTOPTIMIZER", {}),
+                "epochs": ("INT", {"default":1, "min":1, "max":1e6}),
+                "use_gpu": ("BOOLEAN", {"default":False}),
+                "early_stopping": ("BOOLEAN", {"default":False}),
+                "early_stopping_rounds": ("INT", {"default":10, "min":1, "max":1000}),
+                "output_best_val_model": ("BOOLEAN", {"default":True})
+            },
+            "optional": {
+                "val_loader": ("PTDATALOADER", {}),
+            }
+        }
+
+    RETURN_NAMES: tuple = ("Model", "train loss", "val loss")
+    RETURN_TYPES: tuple = ("PTMODEL","TENSOR","TENSOR")
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self, model: nn.Module, 
+          train_loader: torch.utils.data.DataLoader, 
+          optimizer: torch.optim.Optimizer, 
+          epochs: int, 
+          use_gpu: bool, 
+          early_stopping: bool,
+          early_stopping_rounds: int,
+          output_best_val_model:bool,
+          val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
+        """
+        Trains a regression model.
+
+        Args:
+            model (torch.nn.Module): The PyTorch model to be trained.
+            train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.
+            optimizer (torch.optim.Optimizer): Optimizer used for training.
+            epochs (int): Number of training epochs.
+            use_gpu (bool): Whether to use GPU for training.
+            early_stopping (bool): Whether to turn on early stopping
+            early_stopping_rounds (int): 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 (bool): Whether to outputs the model with the best val loss or the model from the last epoch.
+            val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
+
+        Returns:
+            Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, 
+            training loss history, and validation loss history.
+        """
+
+        checkpoint_dir = os.path.realpath(
+            os.path.join(os.path.dirname(__file__), "..", "..", "models", "tmp"))
+        print(f"Checkpoint tmp directory is set to {checkpoint_dir}")
+        if os.path.exists(checkpoint_dir) is False:
+            os.makedirs(checkpoint_dir)
+            print(f"Created {checkpoint_dir}")
+        best_checkpoint_path = os.path.join(checkpoint_dir, "best_model.pt")
+
+        best_val_loss = None
+        rounds_since_best_val_loss = 0
+
+        with torch.inference_mode(False):
+            train_loss = []
+            val_loss = []
+            if use_gpu:
+                model.to("cuda")
+
+            loss_function = nn.MSELoss()  # apply mean loss
+
+            # Train the model
+            for epoch in range(epochs):
+                model.train()
+                total_loss = 0
+                total_samples = len(train_loader.dataset)
+                for images, labels in train_loader:
+                    if use_gpu:
+                        images = images.to("cuda")
+                        labels = labels.to("cuda")
+                    optimizer.zero_grad()
+                    outputs = model(images)
+                    loss = loss_function(outputs, labels)
+                    loss.backward()
+                    optimizer.step()
+ 
+                    batch_size = labels.size(0)
+                    batch_loss = loss.item() * batch_size
+                    total_loss += batch_loss
+
+                epoch_loss = total_loss / total_samples
+                
+                print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
+                train_loss.append(epoch_loss)
+                
+                if val_loader:
+                    with torch.no_grad():
+                        model.train(False)  # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
+                        total_loss = 0
+                        total_samples = len(val_loader.dataset)
+                        for images, labels in val_loader:
+                            if use_gpu:
+                                images = images.to("cuda")
+                                labels = labels.to("cuda")
+                            outputs = model(images)
+                            loss = loss_function(outputs, labels)
+                            batch_size = labels.size(0)
+                            batch_loss = loss.item() * batch_size
+                            total_loss += batch_loss
+
+                        epoch_loss = total_loss / total_samples
+                        print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
+                        val_loss.append(epoch_loss)
+
+                        if best_val_loss is None or epoch_loss < best_val_loss:
+                            best_val_loss = epoch_loss
+                            rounds_since_best_val_loss = 0
+                            torch.save(model.state_dict(), best_checkpoint_path)
+                            print(f"Updated best val loss. Saved the model as {best_checkpoint_path}")
+                        else:
+                            rounds_since_best_val_loss += 1
+                            if early_stopping:
+                                if rounds_since_best_val_loss >= early_stopping_rounds:
+                                    print(f"Val loss did not improve in {early_stopping_rounds} rounds. Stopping training.")
+                                    break
+
+            if use_gpu:
+                model.to("cpu")
+
+            if val_loader:
+                if epoch_loss > best_val_loss and output_best_val_model:
+                    sd = torch.load(best_checkpoint_path, weights_only=True)
+                    model.load_state_dict(sd)
+                    print("Loaded the best model for output")
+                else:
+                    print("Using the final epoch model for output.")
+
+            return (model, 
+                    torch.tensor(train_loss, dtype=torch.float32), 
+                    torch.tensor(val_loss, dtype=torch.float32))

modules/pytorch_wrapper/ptn_linear.py

@@ -0,0 +1,61 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+
+class PtnLinear:
+    """
+    A linear model consisting of a single dense layer.  
+
+    ### Fields:  
+    - `in_features`: The number of input features.
+    - `out_features`: The number of output features.
+    - `bias`: Use bias or not.
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "in_features": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "out_features": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "bias": ("BOOLEAN", {"default": True})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          in_features: int,
+          out_features: int,
+          bias: bool
+          ) -> Tuple[nn.Module]:
+        """
+        Constructs a dense model.
+
+        Args:
+            in_features (int): The number of input features
+            out_features (int): The number of output features
+            bias (bool): Use bias or not.
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = nn.Linear(in_features, out_features, bias)
+            
+        return (model,)

modules/pytorch_wrapper/pto_adam.py

@@ -29,14 +29,14 @@ class PtoAdam:
 
     RETURN_TYPES: tuple = ("PTOPTIMIZER",)
     FUNCTION: str = "f"
-    CATEGORY: str = "Data Analysis"
+    CATEGORY: str = "Training"
 
     def f(self, model, learning_rate, beta1, beta2) -> tuple:
         """
-        Converts PyTorch tensor to NumPy ndarray.
+        Instantiates the optimizer.
 
         Args:
-            tens (torch.Tensor): PyTorch Tensor
+            model (torch.nn.Module): PyTorch model
 
         Returns:
             tuple: A tuple containing the ndarray.

modules/pytorch_wrapper/pto_simple.py

@@ -0,0 +1,88 @@
+from typing import Any, Dict, Union, Iterable
+import torch
+import torch.optim
+
+
+class SimpleOptimizer(torch.optim.Optimizer):
+    """
+    Implementation of the most basic optimizer to update W using the below formula:
+
+$$
+W -= lr * G
+$$
+
+    In Pytorch,
+$$
+param.data -= lr * param.grad
+$$
+
+but accessing param.data directly can be tricky, so we can use:
+$$
+param -= lr * param.grad
+$$
+
+    step() is used instead of forward() so that LR scheduler can call.
+
+    pragma: skip_doc
+    """
+    def __init__(self,
+                 parameters:Iterable[Union[torch.nn.Parameter, Dict[str, Any]]],
+                 lr: float):
+        additional_dict = {"lr": lr}
+        super().__init__(parameters, additional_dict)
+
+    def step(self):
+        with torch.no_grad():
+            # Normally, parameter_group is a list containing a single dict.
+            # However, if the user passes a dict to have separater parameter
+            # for each layer, it can have multiple dicts inside.
+            # param_groups is set up in the super constructor.
+            # You do not need to define this yourself.
+            for g in self.param_groups:  
+                for parameter in g["params"]:
+                    if hasattr(parameter, "grad") and parameter.grad is not None:
+                        parameter -= g["lr"] * parameter.grad
+
+
+class PtoSimple:
+    """
+    Instantiates the most basic optimizer to update W using the below formula:
+$$
+W -= lr * G
+$$
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {}),
+                "learning_rate": ("FLOAT", {"default":0.001, "min":1e-10, "max":1, "step":0.0000000001}),
+            }
+        }
+
+    RETURN_TYPES: tuple = ("PTOPTIMIZER",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    def f(self, model, learning_rate) -> tuple:
+        """
+        Instantiates the optimizer.
+
+        Args:
+            model (torch.nn.Module): PyTorch model
+            learning_rate (float): Learning rate
+
+        Returns:
+            tuple: A tuple containing the ndarray.
+        """
+        with torch.inference_mode(False):
+            opt = SimpleOptimizer(model.parameters(), lr=learning_rate)
+            return(opt,)

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.7"
+version = "3.0.8"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/pt_data_loader_from_tensors_test.py

@@ -0,0 +1,72 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.pt_data_loader_from_tensors import PtDataLoaderFromTensors
+
+class TestPtDataLoader(unittest.TestCase):
+    def setUp(self):
+        """Set up test data."""
+        self.data_loader_node = PtDataLoaderFromTensors()
+
+    def test_with_1d(self):
+        """Tests with transform"""
+        x = torch.Tensor([0, 1, 2, 3])
+        y = torch.Tensor([10, 20, 30, 40])
+
+        bs = 2
+        train_data = self.data_loader_node.f(
+            x=x,
+            y=y,
+            batch_size=bs,
+            shuffle=True
+        )[0]
+
+        x, y = next(iter(train_data))
+        self.assertEqual(x.size(), torch.Size([bs]))
+        self.assertEqual(y.size(), torch.Size([bs]))
+
+    def test_with_2d(self):
+        """Tests with transform"""
+        x = torch.Tensor([0,1,2,3])
+        x = torch.unsqueeze(x, -1)
+        y = torch.Tensor([10,20,30,40])
+
+        bs = 2
+        train_data = self.data_loader_node.f(
+            x=x,
+            y=y,
+            batch_size=bs,
+            shuffle=True
+        )[0]
+
+        x, y = next(iter(train_data))
+        self.assertEqual(x.size(), torch.Size([bs, 1]))
+        self.assertEqual(y.size(), torch.Size([bs]))
+
+    def test_with_2d_no_shuffle(self):
+        """Tests with transform"""
+        x = torch.Tensor([0,1,2,3])
+        x = torch.unsqueeze(x, -1)
+        y = torch.Tensor([10,20,30,40])
+
+        bs = 2
+        train_data = self.data_loader_node.f(
+            x=x,
+            y=y,
+            batch_size=bs,
+            shuffle=False
+        )[0]
+
+        x, y = next(iter(train_data))
+        self.assertEqual(x[0], 0)
+        self.assertEqual(y[0], 10)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/pt_predict_classification_conv_model_test.py

@@ -14,7 +14,7 @@ from pytorch_wrapper.pt_data_loader import PtDataLoader
 from pytorch_wrapper.pt_predict_classification_model import PtPredictClassificationModel
 from pytorch_wrapper.pt_load_model import PtLoadModel
  
-class TestPtPrediction(unittest.TestCase):
+class TestPtPredictionClassificationConvModel(unittest.TestCase):
 
     def setUp(self):
         """Set up test instance."""

tests/pytorch_wrapper/pt_predict_regression_model_test.py

@@ -0,0 +1,44 @@
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.pt_predict_regression_model import PtPredictRegressionModel
+
+
+class TestModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(1, 1, True)
+    def forward(self, inputs):
+        return self.linear(inputs)
+
+class TestPtPredictionRegressionModel(unittest.TestCase):
+
+    def setUp(self):
+        """Set up test instance."""
+        self.model = TestModel()
+        self.x = torch.arange(0, 11, dtype=torch.float32)
+
+    def test_prediction(self):
+        """
+        Tests prediction.  
+
+        Note that this test requires a trained model file on the file system.
+        Unit test will fail if you run this for the first time.
+        Run the unit test for training to generate the weight, and re-run this unit test.
+        """
+        self.predict = PtPredictRegressionModel().f(
+            model = self.model,
+            inputs = self.x,
+            use_gpu = True
+        )[0]
+
+        self.assertTrue(self.predict.size() == torch.unsqueeze(self.x,-1).size())
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptn_linear_model_test.py

@@ -6,7 +6,7 @@ import torch
 PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
 MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
 sys.path.append(MODULE_ROOT)
-from pytorch_wrapper.ptn_linear_model import PtnLinearModel  # Updated import to PtnLinear
+from pytorch_wrapper.ptn_linear_model import PtnLinearModel
 
 
 class TestPtnLinearModel(unittest.TestCase):

tests/pytorch_wrapper/ptn_linear_test.py

@@ -0,0 +1,47 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_linear import PtnLinear
+
+
+class TestPtnLinearModel(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnLinear()
+
+    def test_1(self):
+        self.model = self.model_node.f(
+            2,  # in
+            32,  # out
+            True # bias
+        )[0]
+        x = torch.ones(4, 2)
+        out = self.model(x)
+        self.assertEqual(out.size(), torch.Size([4, 32]))
+        # e contains [0] name and [1] param value, so extract the first element.
+        param_names = list(
+            map(lambda e: e[0], self.model.named_parameters()))
+        self.assertTrue("bias" in param_names)
+
+    def test_2(self):
+        self.model = self.model_node.f(
+            2,  # in
+            1,  # out
+            False # bias
+        )[0]
+        x = torch.ones(4, 2)
+        out = self.model(x)
+        self.assertEqual(out.size(), torch.Size([4, 1]))
+        # e contains [0] name and [1] param value, so extract the first element.
+        param_names = list(
+            map(lambda e: e[0], self.model.named_parameters()))
+        self.assertFalse("bias" in param_names)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/pto_simple_test.py

@@ -0,0 +1,35 @@
+import os
+import sys
+import unittest
+import torch
+
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.pto_simple import PtoSimple
+
+
+class TestPtoSimple(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.node = PtoSimple()
+
+    def test_1(self):
+        model = torch.nn.Linear(2, 3, bias=False)
+        torch.nn.init.constant_(model.weight, 5.0)
+        model.weight.grad = torch.full_like(model.weight, 1.0)
+        opt = self.node.f(model, 0.1)[0]
+        opt.step()  # 5.0 - 0.1 * 1.0 = 4.9
+        expected = 4.9
+        result = model.weight.data[0, 0]
+        self.assertAlmostEqual(
+            result.item(), 
+            expected,
+            places=6, 
+            msg=f"Expected {expected}, got {result}")
+
+
+if __name__ == "__main__":
+    unittest.main()

web/js/ShowTextGeneric.js

@@ -10,7 +10,7 @@ import { ComfyWidgets } from "../../../scripts/widgets.js";
 
 // Displays input text on a node
 app.registerExtension({
-	name: "HowToSD.ShowTextGeneric",
+	name: "HowToSD.PtWrapperShowTextGeneric",
 	async beforeRegisterNodeDef(nodeType, nodeData, app) {
 		if (["PtShowText", "PtShowSize"].includes(nodeData.name)) {  // This needs to match NODE_CLASS_MAPPINGS in Python
 			function populate(text) {