def forward(self, *args):
        """forward pass for calling attention op"""
        assert len(args) == 1
        x = args[0]
        return self.op(x, self.cu_length.tensor())
forward pass for calling attention op
python/aitemplate/frontend/nn/attention.py