Suppose I have a tensor, a
a = torch.Tensor(B,C,X,Y)

I am trying to normalize the Y axis to its Z-score taken along statistics in X. However, some of the elements along the channel dimension, C, are degenerate.
This means I am only interested with those C elements that have a non-zero standard deviation along the axis X.
The way I achieved what I wanted was as follows:
# Find elements with a square-sum standard deviation along the X axis
temp = torch.norm(torch.std(a, axis=-2), dim=-1)
# Find the non-zero indices of this tensor
idx = temp.nonzero()
# Slice the tensor along the indices
a_sliced = a[idx[:,0], idx[:,1], :, :]
# Perform some calculations of a_sliced
a_sliced = foo(a_sliced)
# Slice in the results into a similarly sized tensor as the original tensor.
a_unsliced = torch.zeros_like(a)
a_unsliced[idx[:, 0], idx[:, 1], :, :] = a_unsliced