    def _set_direction(self, direction: Tensor):
        """
        Sets self.max_eig_direction to a normalized version of `direction`
        """
        direction = direction.detach()
        direction = direction / direction.norm()
        direction_sum = direction.sum().item()
        if direction_sum - direction_sum == 0:  # no inf/nan
            self.max_eig_direction[:] = direction
        else:
            logging.info(
                f"Warning: sum of direction in MaxEig is {direction_sum}, "
                "num_channels={self.num_channels}, channel_dim={self.channel_dim}"
            )
Sets self.max_eig_direction to a normalized version of `direction`
modules/general/scaling.py