你好，我在用pytorch_lightning.LightningModule写了一个模型，在传输到SPU时出现error，说模型不是valid JAX array type。我的了解pytorch_lightning.LightningModule是pytorch的wrapper，文档我看到支持pytorch写到XLA。想问为什么会说需要JAX array呢？是不是需要用JAX/FLAX 来写我的模型？
你好，我在用pytorch_lightning.LightningModule写了一个模型，在传输到SPU时出现error，说模型不是valid JAX array type。我的了解pytorch_lightning.LightningModule是pytorch的wrapper，文档我看到支持pytorch写到XLA。想问为什么会说需要JAX array呢？是不是需要用JAX/FLAX 来写我的模型？