pull/1205/head
weishengyu 2021-09-26 14:28:12 +08:00
parent 6529765a0b
commit b84e4352b1
1 changed files with 2 additions and 2 deletions

View File

@ -62,7 +62,7 @@ class PKSampler(DistributedBatchSampler):
elif self.sample_method == "sample_avg_prob":
counter = []
for label_i in self.label_list:
counter.append(len(self.label_list[label_i]))
counter.append(len(self.label_dict[label_i]))
self.prob_list = np.array(counter) / sum(counter)
else:
logger.error(
@ -83,7 +83,7 @@ class PKSampler(DistributedBatchSampler):
np.random.RandomState(self.epoch).shuffle(self.label_list)
for i in range(len(self)):
batch_index = []
batch_label_list = np.random.sample(
batch_label_list = np.random.choice(
self.label_list,
size=label_per_batch,
replace=False,