def METHOD_NAME(self, variables):
    """Initialize optimizer variables.
    SGD optimizer has one variable `momentums`, only set if `self.momentum`
    is not 0.
    Args:
      var_list: list of model variables to build SGD variables on.
    """
    if self.built:
        return
    super().METHOD_NAME(variables)
    self.momentums = []
    if self.momentum != 0:
        for variable in variables:
            self.momentums.append(
                self.add_variable_from_reference(
                    reference_variable=variable, name="momentum"
                )
            )