ErwannMillon / Color-diffusion

A diffusion model to colorize black and white images
631 stars 21 forks source link

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

Open kevio33 opened 1 year ago

kevio33 commented 1 year ago

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 commented 2 months ago

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

kevio33 commented 2 months ago

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