Hi, I am new to PyTorch.
I want to swap elements in a 3D tensor which contains scattered non-zeros elements in shape (batch_size, seq_len, elem_dim), while keeping the gradient tracked. Given an example,
inputs = [[1 0 2]
[0 3 0]
[0 0 4]] # each bold integer represents a vector
outputs = func(inputs)
assert outputs == [[1 2] [3 0] [4 0]]
outputs.sum().backward # raise no errors
how to implement this function?
I come from TensorFlow knowing that this is somehow inconvenient and ineffective to do, but TensorFlow allows gradient tracking with scatter_fn and gather_fn.
If some swap function is provided in newest version of PyTorch, I think I do this.