Program Listing for File batchnorm.h

↰ Return to documentation for file (torch/csrc/api/include/torch/nn/options/batchnorm.h)

#pragma once

#include <torch/arg.h>
#include <torch/csrc/Export.h>
#include <torch/types.h>

namespace torch::nn {

struct TORCH_API BatchNormOptions {
  /* implicit */ BatchNormOptions(int64_t num_features);

  TORCH_ARG(int64_t, num_features);

  TORCH_ARG(double, eps) = 1e-5;

  TORCH_ARG(std::optional<double>, momentum) = 0.1;

  TORCH_ARG(bool, affine) = true;

  TORCH_ARG(bool, track_running_stats) = true;
};

using BatchNorm1dOptions = BatchNormOptions;

using BatchNorm2dOptions = BatchNormOptions;

using BatchNorm3dOptions = BatchNormOptions;

// ============================================================================

namespace functional {

struct TORCH_API BatchNormFuncOptions {
  TORCH_ARG(Tensor, weight) = Tensor();

  TORCH_ARG(Tensor, bias) = Tensor();

  TORCH_ARG(bool, training) = false;

  TORCH_ARG(double, momentum) = 0.1;

  TORCH_ARG(double, eps) = 1e-5;
};

} // namespace functional

} // namespace torch::nn