def METHOD_NAME(self):
    transformer_config = self.transformer_config
    config = transformer_config
    config.recompute_granularity = 'selective'
    selective_transformer_block = TransformerBlock(config)
    assert selective_transformer_block.config.recompute_granularity == 'selective'
    assert selective_transformer_block.checkpoint_core_attention
    sequence_length = 32
    micro_batch_size = 2
    selective_transformer_block.cuda()
    # [sequence length, batch size, hidden size]
    hidden_states = torch.ones((sequence_length, micro_batch_size, config.hidden_size))
    hidden_states = hidden_states.cuda()
    attention_mask = torch.ones((1, 1, sequence_length, sequence_length), dtype=bool).cuda()
    hidden_states = selective_transformer_block(hidden_states=hidden_states, attention_mask=attention_mask)
    assert hidden_states.shape[0] == sequence_length
    assert hidden_states.shape[1] == micro_batch_size
    assert hidden_states.shape[2] == config.hidden_size