Model Works on CPU but Error out while running on GPU

I have a pytorch lightning model where the training is working fine on cpu but when I run the same training on a gpu machine it throws an error on

AssertionError: Gather function not implemented for CPU tensors

What could be the possible issue here.

I use ‘accelerator’: ‘gpu’, ‘strategy’: ‘dp’

The machine has 2 gpus and this is ran on databricks cluster.