You could use a WeightedRandomSampler as described in this post with an example.
If you want to split the dataset in a stratified way, you could use e.g. sklearn.model_selection.train_test_split with the stratify option for the indices of the dataset (and the targets as the inputs for stratify) and use these indices in Subsets.