dropout.h

#pragma once

#include <torch/nn/cloneable.h>
#include <torch/nn/pimpl.h>
#include <torch/types.h>

#include <cstddef>
#include <vector>

namespace torch {
namespace nn {

struct TORCH_API DropoutOptions {
  /* implicit */ DropoutOptions(double rate = 0.5);
  TORCH_ARG(double, rate);
};

namespace detail {
template <typename Derived>
class DropoutImplBase : public torch::nn::Cloneable<Derived> {
 public:
  explicit DropoutImplBase(DropoutOptions options_ = DropoutOptions());

  void reset() override;

  DropoutOptions options;
};
} // namespace detail

class TORCH_API DropoutImpl : public detail::DropoutImplBase<DropoutImpl> {
 public:
  using detail::DropoutImplBase<DropoutImpl>::DropoutImplBase;

  Tensor forward(const Tensor& input);

  void pretty_print(std::ostream& stream) const override;
};

class TORCH_API FeatureDropoutImpl
    : public detail::DropoutImplBase<FeatureDropoutImpl> {
 public:
  using detail::DropoutImplBase<FeatureDropoutImpl>::DropoutImplBase;

  Tensor forward(const Tensor& input);

  void pretty_print(std::ostream& stream) const override;
};

TORCH_MODULE(Dropout);

TORCH_MODULE(FeatureDropout);
} // namespace nn
} // namespace torch