src/transformers/models/beit/modeling_flax_beit.py/FlaxBeitAttention/setup
class FlaxBeitAttention:    def setup(self):
        self.attention = FlaxBeitSelfAttention(self.config, self.window_size, dtype=self.dtype)
        self.output = FlaxBeitSelfOutput(self.config, dtype=self.dtype)