What I need to write in the line level of torch.vison for 21 classes?
I solved my problem. Here are the steps I changed :

add a 'class_name' column to the &quot;marking&quot; DataFrame
replace 'source' with 'class_name' in data splitting (df_folds…)
modify load_image_and_boxes() function to also return labels modify load_cutmix_image_and_boxes() function to also return labels (most work to do here)
replace labels = torch.zeros((boxes.shape[0],), dtype=torch.int64) with torch.tensor(labels, dtype=torch.int64)
in get_net() function, set config.num_classes = number of classes in your dataset