RuntimeError: invalid gradient at index 0 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

1

got error on doing this

from flashtorch.utils import apply_transforms, load_image
from flashtorch.saliency import Backprop
train_dataset = torchvision.datasets.CIFAR10(root='./cifar10/',
                                           train=True,
                                           transform=transforms.Compose([
                                                                         transforms.ToTensor()
                                                                         ]),
                                           download=True)
...
backprop = Backprop(model)
backprop.visualize((train_dataset[0][0].unsqueeze(0)).to('cuda'), 6, guided=True, use_gpu=True)

where do I need to set to(‘cuda’)?