def forward(self, *args):
        """Flattens the input with specified start and end dims."""
        assert len(args) == 1
        x = args[0]
        return self.op(x)
Flattens the input with specified start and end dims.
python/aitemplate/frontend/nn/view_ops.py