kernel error in CUDA when moving Tensors to GPU

Question

I'm trying to replicate this code but I moved my tensors to the GPU, when I run my code on the CPU it runs pretty well, but when I switch to GPU I get this error:

RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

My Code is:

#!/usr/bin/python3

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

if torch.cuda.is_available():
    device_ = torch.device("cuda")
    print("========================\nYou are running on GPU!\n========================")
else:
    device_ = torch.device("cpu")
    print("------------------------\nYou are running on CPU!\n------------------------")

class WineDataset(Dataset):
    
    def __init__(self):
        # data loading
        xy = np.loadtxt("./dataset/wine.csv", dtype=np.float32, delimiter=",",  skiprows=1)
        self.n_samples = xy.shape[0]

        self.x = torch.from_numpy(xy[:, 1:])
        self.x = self.x.to(device_)

        self.y = torch.from_numpy(xy[:, [0]]) # n_samples, 1
        self.y = self.y.to(device_)

        print(self.x.shape, self.y.shape)
        
    def __getitem__(self, index):
        # dataset[0]
        return self.x[index], self.y[index]

    def __len__(self):
        # len(dataset)
        return self.n_samples

dataset = WineDataset()
train_loader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=2)

dataiter = iter(train_loader)
data = next(dataiter)
features, labels = data

I have created this Colab Notebook to make my problem reproducible, can you please tell me how can I run the code with my tensors on the GPU? thanks.

Answer 1

Quoted from pytorch forum: RuntimeError: CUDA error: initialization error "might be raised, if you are trying to initialize the CUDA context multiple times, e.g. if you are using multiple processes"

https://discuss.pytorch.org/t/what-does-runtimeerror-cuda-driver-error-initialization-error-mean/87505

so, you need to change num_workers to 0

train_loader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)