Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[kevio33]

i just run train.py and it log the error below:

File "/content/drive/MyDrive/color-black/Color-diffusion/train.py", line 66, in
trainer.fit(model, train_dl, val_dl)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 529, in fit
call._call_and_handle_interrupt(
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 568, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 973, in _run
results = self._run_stage()
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1016, in _run_stage
self.fit_loop.run()
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
self.advance()
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
self.advance(data_fetcher)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in advance
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/automatic.py", line 178, in run
closure()
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in call
self._result = self.closure(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/automatic.py", line 132, in closure
self._zero_grad_fn()
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/automatic.py", line 217, in zero_grad_fn
self._on_before_zero_grad(optimizer)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/automatic.py", line 281, in _on_before_zero_grad
call._call_lightning_module_hook(trainer, "on_before_zero_grad", optimizer)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 144, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/content/drive/MyDrive/color-black/Color-diffusion/model.py", line 120, in on_before_zero_grad
self.ema.update()
File "/usr/local/lib/python3.10/dist-packages/torch_ema/ema.py", line 117, in update
tmp = (s_param - param)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

what's the problem?

[MagJ1]

I ran into the same problem. I suppose that something is odd in the model.py file.

[kevio33]

sorry, i just saw it. It's been so long and i forget what i did

[YanivDorGalron]

I run into the same problem. fixed it by overriding the .to function of model.

def to(self, device):
     super().to(device)
     self.ema.to(device)