Brilliant  I tried your code and found the time budget is almost affordable now.
Actually when tracing the time cost, I found the most time is spent on torch.split, and creating and assigning the mask.
Thank you~