The torch.square() method returns the square of each element of the input tensor. The output tensor retains the input tensor’s shape and data type. It is compatible with CPUs/GPUs and all numeric data types.
Syntax
torch.square(input, out=None)
Parameters
| Argument | Description |
| input (Tensor) | It is an input tensor of any shape and type. |
| out (Tensor, optional) | It is an output tensor to save the results. |
Element-wise squaring
import torch tensor = torch.tensor([1.0, 10.0, 11.0, 20.0]) print(tensor) # Output: tensor([ 1., 10., 11., 20.]) squared_tensor = torch.square(tensor) print(squared_tensor) # Output: tensor([ 1., 100., 121., 400.])
The output squared_tensor contains the square values of each element of the input tensor.
Multi-Dimensional tensor
If the input tensor is multidimensional, the output tensor will also be multidimensional. Just the values in a tensor will be squared.
import torch tensor_2d = torch.tensor([[1.0, 10.0], [11.0, 20.0]]) print(tensor_2d) # Output: # tensor([[ 1., 10.], # [11., 20.]]) squared_2d_tensor = torch.square(tensor_2d) print(squared_2d_tensor) # Output: # tensor([[ 1., 100.], # [121., 400.]])
Different data types
It does not matter whether the data type is torch.float or torch.int when squaring values. It can handle them really well!import torch # Integer tensor int_tensor = torch.tensor([2, 3, 4], dtype=torch.int32) print(torch.square(int_tensor)) # Output: tensor([4, 9, 16], dtype=torch.int32) # Float tensor float_tensor = torch.tensor([2.5, 3.5], dtype=torch.float32) print(torch.square(float_tensor)) # Output: tensor([6.2500, 12.2500], dtype=torch.float32)
Using the “out” argument
If you have a pre-allocated tensor, you can store the squared element’s result in this tensor.
import torch tensor = torch.tensor([2.0, 3.0, 4.0]) pre_allocated = torch.empty(3) print(pre_allocated) # Output: tensor([0., 0., 0.]) torch.square(tensor, out=pre_allocated) print(pre_allocated) # Output: tensor([ 4., 9., 16.])
The out argument is helpful for memory efficiency in large-scale computations, as it reuses an existing tensor.
Squaring of negative elements
If the input tensor contains negative values, squaring makes them positive.
import torch negative_tensor = torch.tensor([-2.0, -6.0, -8.0]) print(negative_tensor) # Output: tensor([-2., -6., -8.]) positive_tensor = torch.square(negative_tensor) print(positive_tensor) # Output: tensor([ 4., 36., 64.])
The above output suggests that the output tensor contains positive values, as multiplying a negative value by a negative value yields a positive result.
