The torch.arange() method in PyTorch generates a one-dimensional (1D) tensor containing a sequence of values starting from the “start” point (inclusive) and the “end” point (exclusive), incremented by the “step” argument.
You can generate a sequence of indices, coordinate grids, or any scenario where you need a sequence of 1D tensors, you can always use the .arange() method.
Function signature
torch.arange(
start: Number = 0,
end: Number,
step: Number = 1,
*,
out: Optional[Tensor] = None,
dtype: Optional[torch.dtype] = None,
layout: torch.layout = torch.strided,
device: Optional[torch.device] = None,
requires_grad: bool = False
) → Tensor
Arguments
| Name | Value |
| start (Optional) | It is the first value of the sequence. By default, it is zero, but you can set it to any value you want. |
| end | It is the upper bound (exclusive) value of the sequence. |
| step (Optional) | Spacing between the generated values. By default, it is 1. |
| out (Optional) | It is an output tensor in which you write the sequence. It’s helpful when you have a predefined, allocated tensor and you write the sequence to it. |
| dtype (Optional) | It defines the data type of the output tensor. For example, torch.float32, torch.int64, torch.float64, etc.. |
| layout (Optional) | By default, the layout is torch.strided. But you can define the acceptable layouts here. |
| device (Optional) | It is a device on which you place the tensor. (e.g. “cpu” or “cuda:0”). |
| requires_grad (Optional) | By default, it is False, but to enable gradient tracking, you must set it to True. |
Simple Range: Default start and step
Let’s generate a simple 1D tensor sequence.import torch tensor = torch.arange(6) print(tensor) # Output: tensor([0, 1, 2, 3, 4, 5]) print(tensor.shape) # Output: torch.Size([6])
You can see that we quickly generated integer indices for loops or indexing tensors. Only required argument in arange() is “end”, which defines the end integer of the sequence, which is (n-1). So, if the sequence is 6, it will be 0 to (6-1), which is 5.
Specifying both start and end
When you don’t want your sequence to start from 0, you should always define your starting point.
import torch tensor = torch.arange(5, 10) print(tensor) # Output: tensor([5, 6, 7, 8, 9]) print(tensor.shape) # Output: torch.Size(5])
In this code, our starting point is five and up to 10 (10-1 = 9) because 10 is exclusive.
Specifying step
import torch tensor = torch.arange(5, 20, 5) print(tensor) # Output: tensor([ 5, 10, 15]) print(tensor.shape) # Output: torch.Size([3])
Here, you can see that the increment is five between the tensor elements.
Floating-Point sequences
You can generate this kind of floating-point sequence while plotting essential data on a chart.
import torch tensor_float = torch.arange(0, 2, 0.4, dtype=torch.float64) print(tensor_float) # Output: tensor([0.0000, 0.4000, 0.8000, 1.2000, 1.6000], dtype=torch.float64) print(tensor_float.shape) # Output: torch.Size([5])
Specifying device and dtype
Until now, we ran on the CPU device, but you can run on the GPU by passing the “device” argument.
import torch
tensor_float = torch.arange(0, 5,
dtype=torch.float64,
device='cuda:0',
requires_grad=True)
print(tensor_float)
# Output: tensor([0., 1., 2., 3., 4.], dtype=torch.float64, requires_grad=True, #device='cuda:0')
print(tensor_float.shape)
# Output: torch.Size([5])
We also enabled gradient tracking.
In-Place Output
We can create an empty tensor using torch.empty() function to preallocate the memory for the tensor and then write the new sequence on that empty tensor.
import torch out = torch.empty(4, dtype=torch.int64) torch.arange(2, 6, out=out) print(out) # Output: tensor([2, 3, 4, 5]) print(out.shape) # Output: torch.Size([4])This approach is efficient in tight loops. That’s all!
