Question regarding Softmax2d

1

I have the matrix A with dimention of BxCxDxD
I know via softmax 2d i can do apply softmax in channel wise
I was wondering if there is a way to specify how much of the channel to consider.
If i want to use loop i should do it like this:
SoftMax = nn.Softmax2d()
lets say C = 44 and i want to do softmax for first half and second half separately

for i in range(2):
     TempM = A[:,i*22:(i+1)*22,:,:]
    TempM= SoftMax(TempM)
    B_Classes[:,i*22:(i+1)*22,:,:] = TempM

is there any simpler way to do it?