void forward_cpu()

in src/inplace_abn_cpu.cpp [107:142]


void forward_cpu(
    at::Tensor& x_,
    const at::Tensor& mean,
    const at::Tensor& var,
    const c10::optional<at::Tensor>& weight,
    const c10::optional<at::Tensor>& bias,
    float eps,
    Activation activation,
    float activation_param) {
  CHECK_NOT_HALF(x_);

  auto x = normalize_shape(x_);

  // Apply normalization
  auto abs_weight = weight.has_value()
      ? weight.value().abs() + eps
      : at::ones({mean.size(0)}, mean.options());
  auto inv_std = 1 / at::sqrt(var + eps);

  auto scale = weight.has_value() ? abs_weight * inv_std : inv_std;
  auto shift = weight.has_value() ? bias.value() - mean * abs_weight * inv_std
                                  : -mean * inv_std;

  x.mul_(normalize_shape(scale)).add_(normalize_shape(shift));

  switch (activation) {
    case Activation::LeakyReLU:
      at::leaky_relu_(x, activation_param);
      break;
    case Activation::ELU:
      at::elu_(x, activation_param);
      break;
    case Activation::Identity:
      break;
  }
}