I am trying to solve the supervised learning problem (housing price) using pytorch and I encounter CUDA error?

I am getting the following error:

RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_17/1714279870.py in <module>
      4 
      5 
----> 6 data_tensor = CustomData(data)
      7 model = NeuralNetwork().to(device)
      8 train_data_tensor = DataLoader(data_tensor,batch_size,shuffle=True)

/tmp/ipykernel_17/3667328859.py in __init__(self, data, transform, target_transform)
     33         #self.data = data
     34         self.label = data[:,-1]
---> 35         self.label = self.label.to(device)
     36 
     37         self.feature = data[:,:-1]

RuntimeError: CUDA error: device-side assert triggered
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.

My code is

import numpy as np 
import pandas as pd 
import matplotlib.pyplot as plt
from pandas import DataFrame
from torch.utils.data import Dataset
import torch

import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
import torch.optim as optim
from torch.autograd import Variable
import torch.nn.functional as F
import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

Here I try fill gaps in the data and also convert text data (like good, and bad) into numbers and finally convert the entire thing into a tensor. So far no error.

data = pd.read_csv('../input/house-prices-advanced-regression-techniques/train.csv')
data_cat = data[[str(i) for i in data.columns if data[str(i)].dtype == object]]
data_cat = data_cat.rank(axis=0, method='dense', na_option='top', ascending=False).astype(float)#converting text data to numbers
data_cat.fillna(0)
data_num = data[[str(i) for i in data.columns if data[str(i)].dtype != object]]
data_val = data_num.iloc[:,-1]
data_val = data_val.fillna(0)
data_num = data_num.iloc[:,:-1]
data_num = (data_num - data_num.mean())/(data_num.std())
data_num = data_num.fillna(0)
data = pd.concat([data_cat,data_num,data_val],axis=1)
data = torch.from_numpy(data.to_numpy())

Now I create a dataset to be used in NeuralNetwork via Dataloader. Again so far so good.

class CustomDataset(Dataset):
    def __init__(self,feature, label, transform=None, target_transform=None):
        self.feature = feature
        self.label = label
        self.transform = transform
        self.target_transform = target_transform
    def __len__(self):
        return len(self.label)
    def __getitem__(self, idx):
        feature = self.feature[idx]
        label = self.label[idx]
        if self.transform:
            feature = self.transform(feature)
        if self.target_transform:
            label = self.target_transform(label)
        return feature, label      
data_tensor = CustomDataset(data[:,:-1],data[:,-1])

Since I will be using GPU, I define the device.

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

This is my nn model. Because there are categorical data, I am using an embedding layer first.

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        
        self.linear_relu_stack = nn.Sequential(
            nn.Embedding(79,10),
            nn.Linear(70, 10),
            nn.ReLU(),
            nn.Linear(10, 10),
            nn.ReLU(),
            nn.Linear(10, 10),
            nn.ReLU(),            
            nn.Linear(10, 1),
        )
    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits

These are my hyperprameters.

learning_rate = 0.005
batch_size = 20
num_of_epoch = 50

Now I am setting up the training process. I suspect this is where the bug lies.

model = NeuralNetwork().to(device)
train_data_tensor = DataLoader(data_tensor,batch_size,shuffle=True)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
def train_loop(data_loader, model, loss_fn, optimizer):
    size = len(data_loader.dataset)
    for batch, (X, y) in enumerate(data_loader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)
        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()        
        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

#Actual training
for t in range(num_of_epoch):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_data_tensor, model, loss_fn, optimizer)
    
print("Done")

Can someone help me rectify errors in code?