Make nn.Square/nn.Threshold handle non-contiguous inputs