torch.nn.Softmax3d?

We have the “SoftMax over features to each spatial location.” for 2D spatial locations, i.e.,the input is (N,C,H,W). And the softmax2d applies Softmax to each 2d location (Channels, h_i, w_j).

I didn’t find SoftMax3d. So I am wondering would there be SoftMax3d in the near future? Like, the input is (N,C,D,H,W), and the softmax3d applies Softmax to each 3d location (Channels, d_k, h_i, w_j).

1 Like