kan/tutorials/API_10_device.ipynb errors out, tensors are not on the same device

[hrshtt]

the tutorial example API_10_device errors out with the below stack trace

description:   0%|                                                           | 0/50 [00:01<?, ?it/s]
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-7-3f2157aed2aa>](https://localhost:8080/#) in <cell line: 7>()
      5 # train the model
      6 #model.train(dataset, opt="LBFGS", steps=20, lamb=1e-3, lamb_entropy=2.);
----> 7 model.train(dataset, opt="LBFGS", steps=50, lamb=5e-5, lamb_entropy=2.);

6 frames
[/usr/local/lib/python3.10/dist-packages/torch/functional.py](https://localhost:8080/#) in einsum(*args)
    378         # the path for contracting 0 or 1 time(s) is already optimized
    379         # or the user has disabled using opt_einsum
--> 380         return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
    381 
    382     path = None

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

reproduced on a fresh install on colab and local system

[sourenaKhanzadeh]

Im getting a similar error on my project, both model and dataset are on the same device, for me it stops working when you forward run the model with train_inputs, like: model(dataset['train_input']), maybe try device=device in the train method as well and see if it is resolved, most likely this is due to the fact that the torch versions are different.