In ch01-38, how does "x: torch.Tensor" in "def forward(self, x: torch.Tensor)" work?

[Reetos]

I have never seen an argument is written like this for a function.
What kind of input is this?

# Create a Linear Regression model class
class LinearRegressionModel(nn.Module): # <- almost everything in PyTorch is a nn.Module (think of this as neural network lego blocks)
    def __init__(self):
        super().__init__() 
        self.weights = nn.Parameter(torch.randn(1, # <- start with random weights (this will get adjusted as the model learns)
                                                dtype=torch.float), # <- PyTorch loves float32 by default
                                   requires_grad=True) # <- can we update this value with gradient descent?)

        self.bias = nn.Parameter(torch.randn(1, # <- start with random bias (this will get adjusted as the model learns)
                                            dtype=torch.float), # <- PyTorch loves float32 by default
                                requires_grad=True) # <- can we update this value with gradient descent?))

    # Forward defines the computation in the model
    def forward(self, x: torch.Tensor) -> torch.Tensor: # <- "x" is the input data (e.g. training/testing features)
        return self.weights * x + self.bias # <- this is the linear regression formula (y = m*x + b)

[mrdbourke]

Hi @Reetos,

Good question!

I answered a similar question here: #360 (comment)

But I'll paste it here too (because the question is similar):

Those are called "type annotations" or "type hints" in Python.

def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.weights * x + self.bias

The x: torch.Tensor means that the forward() method is expecting x to be of type torch.Tensor.

And the -> torch.Tensor means that the forward() method returns a torch.Tensor too.

Type hints/annotations are helpful for when you'd like to know the inputs/outputs of your code/functions.

They don't necessarily return errors but when looking up documentation for what code does, the type annotation/hint shows you what the code is expecting as input/output.

You can see more here: https://realpython.com/python-type-checking/

Or here: https://docs.python.org/3/library/typing.html

[Reetos]

I got it.

Thank you for your answer!

[avish121]

your question helped me too.

[Mingfei-Jiang]

great answer!