Does pytorch models require particular kind of class labels?

I am trying to take a subset of my dataset using Pandas .sample method and due to the highly unbalanced structure of my data’s class distribution, I lose some of the classes which is ok. However, when I want to train the model I run into this:

RuntimeError: cuda runtime error (710) : device-side assert triggered at

I have run into this error in the past, and it always meant that something was wrong with my model architecture especially the last layer. Is this error showing up because of the lost classes during sampling?

# sampling
sample_csv = train_csv.sample(frac = .2)
sample_csv.reset_index(drop=True, inplace = True)
sample_landmark_count = sample_csv.landmark_id.value_counts()
# resnet model
num_classes = len(sample_landmark_count) 
resnet = models.resnet18()
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, num_classes)    

The labels are in the landmark_id and the data are from kaggle’s google landmark competition.

And this is how I made my costum dataset

class landmark_dataset(Dataset):
   def __init__(self,csv_file,mode , tfms = None):
       self.mode = mode
       self.csv_file = csv_file
       self.transforms = tfms
       
   def __len__(self):
       
       return len(self.csv_file)
       
   def __getitem__(self, idx):
       
       img_path = get_folder_names(self.csv_file.id[idx])
       img = Image.open(img_path)
       img = self.transforms(img)
       if self.mode == 'train':
           labels = self.csv_file.landmark_id[idx]
           return img, labels 
       else:
           return img

Which PyTorch version are you using? If not 1.6, could you update and rerun the code?
Also, could you run the code via CUDA_LAUNCH_BLOCKING=1 python script.py args and post the stack trace here?