def forward()

in models/quantize_affine.py [0:0]

• 32 lines of code
• 9 McCabe index (conditional complexity)

    def forward(self, input, recomp_bn_stats=False, override_alpha=False):
        if (
            self.training
            and self.is_active()
            and not self.quantizer_freeze_min_max
        ):
            # Force include 0 in min_value and max_value calculation.
            min_value = min(input.min().item(), 0)
            max_value = max(input.max().item(), 0)

            if self.iteration_count == self.iteration_delay:
                new_running_min_value = min_value
                new_running_max_value = max_value
            else:
                new_running_min_value = (
                    1.0 - self.momentum
                ) * self.running_min_value.item() + self.momentum * min_value
                new_running_max_value = (
                    1.0 - self.momentum
                ) * self.running_max_value.item() + self.momentum * max_value

            self.running_min_value.fill_(new_running_min_value)
            self.running_max_value.fill_(new_running_max_value)

        if self.is_active():
            output = quantize_affine(
                input,
                self.running_min_value.item(),
                self.running_max_value.item(),
                self.num_bits,
            )
        else:
            output = input

        if self.training and self.increment_counter:
            self.iteration_count.fill_(self.iteration_count.item() + 1)

        return output