def update(self, x: Tensor) -> None:
        """Update the center with a new batch of features.

        Args:
            x:
                Feature tensor used to update the center. Must have the same number of
                dimensions as self.size.
        """
        batch_center = self._center_fn(x=x, dim=self.dim)
        self.center = center_momentum(
            center=self.center, batch_center=batch_center, momentum=self.momentum
        )
Update the center with a new batch of features.

        Args:
            x:
                Feature tensor used to update the center. Must have the same number of
                dimensions as self.size.
        
lightly/models/modules/center.py