aten/contrib/data/ShuffleDataset.cc


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

#include "ShuffleDataset.h"
#include "Dataset.h"
#include <algorithm>

using namespace at;

ShuffleDataset::ShuffleDataset(Dataset& dataset) : ResampleDataset(dataset) {
   resample();
}

void ShuffleDataset::resample() {
   perm_.reserve(size_);
   for(int n = 0; n < size_; ++n)
      perm_[n] = n;
   std::random_shuffle(perm_.begin(), perm_.end());
}