RuntimeError: expected device cuda:0 but got device cpu

[srmykola]

Hi Sachin,
Thanks for moving the prophet idea to GPUs.
While using your library I am running into the issue:

0.00% [0/30 00:00<00:00]

0.00% [0/22 00:00<00:00]

RuntimeError Traceback (most recent call last)
in
4 # , alpha=1e-2
5 )
----> 6 model.fit(df)

~/profetorch-test/env/lib/python3.7/site-packages/profetorch/model/model.py in fit(self, df)
63 # breakpoint()
64
---> 65 learner.fit_one_cycle(self.epochs, self.lr, callbacks=cb)
66
67 self.model = learner.model

~/profetorch-test/env/lib/python3.7/site-packages/fastai/train.py in fit_one_cycle(learn, cyc_len, max_lr, moms, div_factor, pct_start, final_div, wd, callbacks, tot_epochs, start_epoch)
21 callbacks.append(OneCycleScheduler(learn, max_lr, moms=moms, div_factor=div_factor, pct_start=pct_start,
22 final_div=final_div, tot_epochs=tot_epochs, start_epoch=start_epoch))
---> 23 learn.fit(cyc_len, max_lr, wd=wd, callbacks=callbacks)
24
25 def fit_fc(learn:Learner, tot_epochs:int=1, lr:float=defaults.lr, moms:Tuple[float,float]=(0.95,0.85), start_pct:float=0.72,

~/profetorch-test/env/lib/python3.7/site-packages/fastai/basic_train.py in fit(self, epochs, lr, wd, callbacks)
198 else: self.opt.lr,self.opt.wd = lr,wd
199 callbacks = [cb(self) for cb in self.callback_fns + listify(defaults.extra_callback_fns)] + listify(callbacks)
--> 200 fit(epochs, self, metrics=self.metrics, callbacks=self.callbacks+callbacks)
201
202 def create_opt(self, lr:Floats, wd:Floats=0.)->None:

~/profetorch-test/env/lib/python3.7/site-packages/fastai/basic_train.py in fit(epochs, learn, callbacks, metrics)
99 for xb,yb in progress_bar(learn.data.train_dl, parent=pbar):
100 xb, yb = cb_handler.on_batch_begin(xb, yb)
--> 101 loss = loss_batch(learn.model, xb, yb, learn.loss_func, learn.opt, cb_handler)
102 if cb_handler.on_batch_end(loss): break
103

~/profetorch-test/env/lib/python3.7/site-packages/fastai/basic_train.py in loss_batch(model, xb, yb, loss_func, opt, cb_handler)
24 if not is_listy(xb): xb = [xb]
25 if not is_listy(yb): yb = [yb]
---> 26 out = model(*xb)
27 out = cb_handler.on_loss_begin(out)
28

~/profetorch-test/env/lib/python3.7/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
548 result = self._slow_forward(*input, **kwargs)
549 else:
--> 550 result = self.forward(*input, **kwargs)
551 for hook in self._forward_hooks.values():
552 hook_result = hook(self, input, result)

~/profetorch-test/env/lib/python3.7/site-packages/profetorch/model/blocks.py in forward(self, t, x)
236 prediction = torch.cat([m(t,x) for m in self.models], -1)
237 median = prediction[:, [self.idx]]
--> 238 prediction[:, self.idxs] = median + F.softplus(prediction[:, self.idxs]) * self.signs
239
240 return self.squash(prediction)

RuntimeError: expected device cuda:0 but got device cpu

OS: ubuntu
NVIDIA-SMI 418.87.01 Driver Version: 418.87.01 CUDA Version: 10.1
torch==1.5.0+cu101

Let me know what could be done.
Thank you!