p_in Class — pytorch Architecture

Source Code

aten/src/ATen/core/DistributionsHelper.h lines 212–227

template <typename T>
struct bernoulli_distribution {

  C10_HOST_DEVICE inline bernoulli_distribution(T p_in) : p(p_in) {
    TORCH_CHECK_IF_NOT_ON_CUDA(p_in >= 0 && p_in <= 1);
  }

  template <typename RNG>
  C10_HOST_DEVICE inline T operator()(RNG* generator) const {
    uniform_real_distribution<T> uniform(0.0, 1.0);
    return transformation::bernoulli<T>(uniform(generator), p);
  }

  private:
    T p;
};