Yeah sure.
class_sample_count = np.array([len(np.where(Ytr == t)[0]) for t in np.unique(Ytr)])
weight = 1. / class_sample_count
samples_weight = np.array([weight[t] for t in Ytr])
samples_weight = torch.from_numpy(samples_weight)
sampler = torch.utils.data.sampler.WeightedRandomSampler(samples_weight.type('torch.DoubleTensor'), len(samples_weight))
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, num_workers=1, sampler=sampler)

Maybe I will try to under sample negative sets and try to see how it behaves.