Learn more  » Push, build, and install  RubyGems npm packages Python packages Maven artifacts PHP packages Go Modules Bower components Debian packages RPM packages NuGet packages

neilisaac / torch   python

Repository URL to install this package:

Version: 1.8.0 

/ include / caffe2 / operators / weighted_multi_sampling_op.h

#pragma once
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"

namespace caffe2 {

template <class Context>
class WeightedMultiSamplingOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

  template <class... Args>
  explicit WeightedMultiSamplingOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        num_samples_(
            this->template GetSingleArgument<int64_t>("num_samples", 0)) {
    CAFFE_ENFORCE_GE(num_samples_, 0);
  }

  bool RunOnDevice() override;

 private:
  const int64_t num_samples_;
};

} // namespace caffe2