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 trainer.fit(model)
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.