Working GPU Example for demo_bnn_classification.ipynb

[janfelixklein]

Hallo guys!

I just wanted to try out your example as mentioned in header. However I get the error

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

when running the fit method of the BnnClassification. Have you tried this example on a GPU machine? When I set device = torch.device("cpu"), it runs smoothly.

I have taken a closer look in your implementions and found several places, where tensors are created without having the device in mind.

Best regards

JFK