Select input slices to calculate gradient for with torch.autograd.grad

The error is expected since a new variable is created through the differentiable slice operation. A workaround is described here.