How should I solve "Thread group count is above the maximum allowed limit" error while running AI model?

Would you be able to suggest how to tackle the issue getting Thread group count is above the maximum allowed limit error while running an AI model?

The exact error I get is
Exceeded safe compute dispatch group count limit per dimension [1, 127008, 1] for ScatterOps (UnityEngine.ComputeShader)
Thread group count is above the maximum allowed limit. Maximum allowed thread group count is 65535.

AI model gives reasonable output even with the error but slightly different from an expected output.

Should I reduce the input size? Currently its (1, 1, 256, 256) tensor from image.
Or is there any other potential solution you can suggest?
Maybe modifying something around the Scatter node?
Thank you for help!

Yes, the recommended workaround is to reduce the input size.
As of modifying the model, it is hard to tell. If you share the original model file, we can take a look into it.