请问deepONet_POD如何设置多卡训练

[Qinger27]

你好！
我在40G的A100卡上训练您的burgers方程的deepONet.py文件时，报错oom：

请问下怎么修改来完成多卡训练呀？抱歉我对DeepXDE的使用不是很熟悉，期待回复~

[lululxvi]

The demo code shouldn't use so much memory.