def forward(self, *args):
        """Applies Softmax on the input tensor."""
        assert len(args) == 1
        x = args[0]

        return ops.softmax(x, self.dim)
Applies Softmax on the input tensor.
python/aitemplate/frontend/nn/softmax.py