lenet5_pytorch.ipynb: Executing code from "In [12]" crashes program

The line ```_, probs = model(valid_dataset[index][0].unsqueeze(0))``` throws ```RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same```