2D input in Deeponet

[xyyrly]

Hi all,

I am trying to train a Deeponet for a diffusion equation with a source term, s(x,t). My question is: should I flatten the 2d source term samples and feed them into my branch net, just like 1D inputs? Or is there any open-sourced examples with 2d input? Thank you very much!

Best,
Xiaoyu