torch.Tensor.index_fill_
torch.Tensor.index_fill_(dim, index, value) → Tensor
Fills the elements of the self
tensor with value value
by selecting the indices in the order given in index
.
- Parameters
-
-
dim (int) – dimension along which to index
-
index (LongTensor) – indices of
self
tensor to fill in -
value (float) – the value to fill with
-
Example::
>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) >>> index = torch.tensor([0, 2]) >>> x.index_fill_(1, index, -1) tensor([[-1., 2., -1.], [-1., 5., -1.], [-1., 8., -1.]])
再看一个例子:
a = torch.zeros((2, 5)) index = torch.LongTensor([0, 2]) a.index_fill(1, index, 1) >>> tensor([[1., 0., 1., 0., 0.], [1., 0., 1., 0., 0.]])