DeepSpeed: RuntimeError: expected scalar type Float but found Half

Hi - I’m trying to use the deepspeed plugin with Pytorch Lightning. My code worked before but changing the line in trainer

to add

plugins='deepspeed_stage_3_offload'

Causes the error posted in the title. I’ve tried casting parameters and variables as float and half, but the error persists.

Any suggestions would be much appreciated as I’m really looking forward to see what deepspeed can do.

I should note that the error is happening in a call to a pytorch_geometric method (if that changes anything).

deepspeed==0.4.3 pytorch-lightning==1.3.8 torch==1.9.0 torch-cluster==1.5.9 torch-geometric==1.7.1 torch-scatter==2.0.7 torch-sparse==0.6.10 torch-spline-conv==1.2.1 torchmetrics==0.3.2 torchvision==0.10.0