Thanks for reaching out Tanmay and glad you were able to get attributions working with WIT in notebook mode.
Unfortunately, WIT in TensorBoard doesn't currently support attributions. This is because WIT in TensorBoard queries the models through TF Serving's PredictionService API, which only returns prediction results from the served model, and not other custom calculated values such as shap values. We rely on the model being queried to provide the attribution results (just as you did with a custom prediction fn when using attributions in notebook mode).
When you mention wanting to use TB mode due to a bigger data size, what exactly do you refer to? Are you able to load more data points when using WIT in TB vs in notebook mode, or were you just experimenting with the limits of each mode?