Here is the code:
class_sample_count = torch.tensor(
[(torch.tensor(train_dataset.targets) == t).sum()
for t in torch.unique(torch.tensor([train_dataset.targets]), sorted=True)])
weight = torch.div(1., class_sample_count.float())
train_samples_weight = torch.tensor([weight[t.item()] for t in torch.tensor(train_dataset.targets)[train_idx]])
train_sampler = torch.utils.data.WeightedRandomSampler(train_samples_weight, len(train_samples_weight))